Include PRS too

Signed-off-by: Tim Vaillancourt <[email protected]>

go/cmd/vtctldclient/command/reparents.go

@@ -195,6 +195,7 @@ func commandInitShardPrimary(cmd *cobra.Command, args []string) error {
 var plannedReparentShardOptions = struct {
 	NewPrimaryAliasStr      string
 	AvoidPrimaryAliasStr    string
+	ExpectedPrimaryAliasStr string
 	WaitReplicasTimeout     time.Duration
 	TolerableReplicationLag time.Duration
 	AllowCrossCellPromotion bool
@@ -207,8 +208,9 @@ func commandPlannedReparentShard(cmd *cobra.Command, args []string) error {
 	}
 
 	var (
-		newPrimaryAlias   *topodatapb.TabletAlias
-		avoidPrimaryAlias *topodatapb.TabletAlias
+		newPrimaryAlias      *topodatapb.TabletAlias
+		avoidPrimaryAlias    *topodatapb.TabletAlias
+		expectedPrimaryAlias *topodatapb.TabletAlias
 	)
 
 	if plannedReparentShardOptions.NewPrimaryAliasStr != "" {
@@ -225,13 +227,21 @@ func commandPlannedReparentShard(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	if plannedReparentShardOptions.ExpectedPrimaryAliasStr != "" {
+		expectedPrimaryAlias, err = topoproto.ParseTabletAlias(plannedReparentShardOptions.ExpectedPrimaryAliasStr)
+		if err != nil {
+			return err
+		}
+	}
+
 	cli.FinishedParsing(cmd)
 
 	resp, err := client.PlannedReparentShard(commandCtx, &vtctldatapb.PlannedReparentShardRequest{
 		Keyspace:                keyspace,
 		Shard:                   shard,
 		NewPrimary:              newPrimaryAlias,
 		AvoidPrimary:            avoidPrimaryAlias,
+		ExpectedPrimary:         expectedPrimaryAlias,
 		WaitReplicasTimeout:     protoutil.DurationToProto(plannedReparentShardOptions.WaitReplicasTimeout),
 		TolerableReplicationLag: protoutil.DurationToProto(plannedReparentShardOptions.TolerableReplicationLag),
 		AllowCrossCellPromotion: plannedReparentShardOptions.AllowCrossCellPromotion,
@@ -310,6 +320,7 @@ func init() {
 	PlannedReparentShard.Flags().DurationVar(&plannedReparentShardOptions.TolerableReplicationLag, "tolerable-replication-lag", 0, "Amount of replication lag that is considered acceptable for a tablet to be eligible for promotion when Vitess makes the choice of a new primary.")
 	PlannedReparentShard.Flags().StringVar(&plannedReparentShardOptions.NewPrimaryAliasStr, "new-primary", "", "Alias of a tablet that should be the new primary.")
 	PlannedReparentShard.Flags().StringVar(&plannedReparentShardOptions.AvoidPrimaryAliasStr, "avoid-primary", "", "Alias of a tablet that should not be the primary; i.e. \"reparent to any other tablet if this one is the primary\".")
+	PlannedReparentShard.Flags().StringVar(&plannedReparentShardOptions.NewPrimaryAliasStr, "expected-primary", "", "Alias of a tablet that must be the current primary in order for the reparent to be processed.")
 	PlannedReparentShard.Flags().BoolVar(&plannedReparentShardOptions.AllowCrossCellPromotion, "allow-cross-cell-promotion", false, "Allow cross cell promotion")
 	Root.AddCommand(PlannedReparentShard)
 

go/vt/vtctl/grpcvtctldserver/server.go

@@ -3072,6 +3072,10 @@ func (s *VtctldServer) PlannedReparentShard(ctx context.Context, req *vtctldatap
 		span.Annotate("avoid_primary_alias", topoproto.TabletAliasString(req.AvoidPrimary))
 	}
 
+	if req.ExpectedPrimary != nil {
+		span.Annotate("expected_primary_alias", topoproto.TabletAliasString(req.ExpectedPrimary))
+	}
+
 	if req.NewPrimary != nil {
 		span.Annotate("new_primary_alias", topoproto.TabletAliasString(req.NewPrimary))
 	}
@@ -3091,6 +3095,7 @@ func (s *VtctldServer) PlannedReparentShard(ctx context.Context, req *vtctldatap
 		reparentutil.PlannedReparentOptions{
 			AvoidPrimaryAlias:       req.AvoidPrimary,
 			NewPrimaryAlias:         req.NewPrimary,
+			ExpectedPrimaryAlias:    req.ExpectedPrimary,
 			WaitReplicasTimeout:     waitReplicasTimeout,
 			TolerableReplLag:        tolerableReplLag,
 			AllowCrossCellPromotion: req.AllowCrossCellPromotion,

go/vt/vtctl/reparentutil/planned_reparenter.go

@@ -61,6 +61,7 @@ type PlannedReparenter struct {
 type PlannedReparentOptions struct {
 	NewPrimaryAlias         *topodatapb.TabletAlias
 	AvoidPrimaryAlias       *topodatapb.TabletAlias
+	ExpectedPrimaryAlias    *topodatapb.TabletAlias
 	WaitReplicasTimeout     time.Duration
 	TolerableReplLag        time.Duration
 	AllowCrossCellPromotion bool
@@ -507,6 +508,13 @@ func (pr *PlannedReparenter) reparentShardLocked(
 		return err
 	}
 
+	if opts.ExpectedPrimaryAlias != nil && !topoproto.TabletAliasEqual(opts.ExpectedPrimaryAlias, shardInfo.PrimaryAlias) {
+		return vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "primary %s is not equal to expected alias %s",
+			topoproto.TabletAliasString(shardInfo.PrimaryAlias),
+			topoproto.TabletAliasString(opts.ExpectedPrimaryAlias),
+		)
+	}
+
 	keyspaceDurability, err := pr.ts.GetKeyspaceDurability(ctx, keyspace)
 	if err != nil {
 		return err

go/vt/vtctl/reparentutil/planned_reparenter_flaky_test.go

@@ -2510,6 +2510,7 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) {
 		name       string
 		tmc        tmclient.TabletManagerClient
 		tablets    []*topodatapb.Tablet
+		shards     []*vtctldatapb.Shard
 		unlockTopo bool
 
 		ev       *events.Reparent
@@ -3296,6 +3297,73 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "expected primary mismatch",
+			tmc: &testutil.TabletManagerClient{
+				GetGlobalStatusVarsResults: map[string]struct {
+					Statuses map[string]string
+					Error    error
+				}{
+					"zone1-0000000200": {
+						Statuses: map[string]string{
+							InnodbBufferPoolsDataVar: "123",
+						},
+					},
+					"zone1-0000000100": {
+						Statuses: map[string]string{
+							InnodbBufferPoolsDataVar: "123",
+						},
+					},
+				},
+			},
+			shards: []*vtctldatapb.Shard{
+				{
+					Keyspace: "testkeyspace",
+					Name:     "-",
+					Shard: &topodatapb.Shard{
+						IsPrimaryServing: true,
+						PrimaryAlias: &topodatapb.TabletAlias{
+							Cell: "zone1",
+							Uid:  100,
+						},
+					},
+				},
+			},
+			tablets: []*topodatapb.Tablet{
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "zone1",
+						Uid:  100,
+					},
+					Type:     topodatapb.TabletType_PRIMARY,
+					Keyspace: "testkeyspace",
+					Shard:    "-",
+				},
+				{
+					Alias: &topodatapb.TabletAlias{
+						Cell: "zone1",
+						Uid:  200,
+					},
+					Type:     topodatapb.TabletType_REPLICA,
+					Keyspace: "testkeyspace",
+					Shard:    "-",
+				},
+			},
+
+			ev:       &events.Reparent{},
+			keyspace: "testkeyspace",
+			shard:    "-",
+			opts: PlannedReparentOptions{
+				// This is not the shard primary, so it should cause an error.
+				ExpectedPrimaryAlias: &topodatapb.TabletAlias{
+					Cell: "zone1",
+					Uid:  200,
+				},
+			},
+
+			shouldErr:     true,
+			expectedEvent: nil,
+		},
 	}
 
 	logger := logutil.NewMemoryLogger()
@@ -3312,9 +3380,13 @@ func TestPlannedReparenter_reparentShardLocked(t *testing.T) {
 			testutil.AddTablets(ctx, t, ts, &testutil.AddTabletOptions{
 				AlsoSetShardPrimary:  true,
 				ForceSetShardPrimary: true, // Some of our test cases count on having multiple primaries, so let the last one "win".
-				SkipShardCreation:    false,
+				SkipShardCreation:    len(tt.shards) > 0,
 			}, tt.tablets...)
 
+			if len(tt.shards) > 0 {
+				testutil.AddShards(ctx, t, ts, tt.shards...)
+			}
+
 			if !tt.unlockTopo {
 				lctx, unlock, err := ts.LockShard(ctx, tt.keyspace, tt.shard, "locking for testing")
 				require.NoError(t, err, "could not lock %s/%s for testing", tt.keyspace, tt.shard)

proto/vtctldata.proto

@@ -1386,6 +1386,9 @@ message PlannedReparentShardRequest {
   vttime.Duration tolerable_replication_lag = 6;
   // AllowCrossCellPromotion allows cross cell promotion,
   bool allow_cross_cell_promotion = 7;
+  // ExpectedPrimary is the optional alias we expect to be the current primary in order for
+  // the reparent operation to succeed.
+  topodata.TabletAlias expected_primary = 8;
 }
 
 message PlannedReparentShardResponse {