kvserver: modernize AdminRelocateRange

pkg/kv/batch.go

@@ -742,7 +742,9 @@ func (b *Batch) adminChangeReplicas(
 // adminRelocateRange is only exported on DB. It is here for symmetry with the
 // other operations.
 func (b *Batch) adminRelocateRange(
-	key interface{}, voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	key interface{},
+	voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	transferLeaseToFirstVoter bool,
 ) {
 	k, err := marshalKey(key)
 	if err != nil {
@@ -753,8 +755,10 @@ func (b *Batch) adminRelocateRange(
 		RequestHeader: roachpb.RequestHeader{
 			Key: k,
 		},
-		VoterTargets:    voterTargets,
-		NonVoterTargets: nonVoterTargets,
+		VoterTargets:                      voterTargets,
+		NonVoterTargets:                   nonVoterTargets,
+		TransferLeaseToFirstVoter:         transferLeaseToFirstVoter,
+		TransferLeaseToFirstVoterAccurate: true,
 	}
 	b.appendReqs(req)
 	b.initResult(1, 0, notRaw, nil)

pkg/kv/db.go

@@ -655,10 +655,13 @@ func (db *DB) AdminChangeReplicas(
 // AdminRelocateRange relocates the replicas for a range onto the specified
 // list of stores.
 func (db *DB) AdminRelocateRange(
-	ctx context.Context, key interface{}, voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	ctx context.Context,
+	key interface{},
+	voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	transferLeaseToFirstVoter bool,
 ) error {
 	b := &Batch{}
-	b.adminRelocateRange(key, voterTargets, nonVoterTargets)
+	b.adminRelocateRange(key, voterTargets, nonVoterTargets, transferLeaseToFirstVoter)
 	return getOneErr(db.Run(ctx, b), b)
 }
 

pkg/kv/kvserver/client_relocate_range_test.go

@@ -47,7 +47,11 @@ func relocateAndCheck(
 	testutils.SucceedsSoon(t, func() error {
 		err := tc.Servers[0].DB().
 			AdminRelocateRange(
-				context.Background(), startKey.AsRawKey(), voterTargets, nonVoterTargets,
+				context.Background(),
+				startKey.AsRawKey(),
+				voterTargets,
+				nonVoterTargets,
+				true, /* transferLeaseToFirstVoter */
 			)
 		if err != nil {
 			if every.ShouldLog() {
@@ -77,7 +81,11 @@ func requireRelocationFailure(
 ) {
 	testutils.SucceedsSoon(t, func() error {
 		err := tc.Servers[0].DB().AdminRelocateRange(
-			ctx, startKey.AsRawKey(), voterTargets, nonVoterTargets,
+			ctx,
+			startKey.AsRawKey(),
+			voterTargets,
+			nonVoterTargets,
+			true, /* transferLeaseToFirstVoter */
 		)
 		if kv.IsExpectedRelocateError(err) {
 			return err
@@ -299,6 +307,44 @@ func TestAdminRelocateRange(t *testing.T) {
 	}
 }
 
+// TestAdminRelocateRangeWithoutLeaseTransfer tests that `AdminRelocateRange`
+// only transfers the lease away to the first voting replica in the target slice
+// if the callers asks it to.
+func TestAdminRelocateRangeWithoutLeaseTransfer(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	args := base.TestClusterArgs{
+		ReplicationMode: base.ReplicationManual,
+	}
+
+	tc := testcluster.StartTestCluster(t, 5 /* numNodes */, args)
+	defer tc.Stopper().Stop(ctx)
+
+	k := keys.MustAddr(tc.ScratchRange(t))
+
+	// Add voters to the first three nodes.
+	relocateAndCheck(t, tc, k, tc.Targets(0, 1, 2), nil /* nonVoterTargets */)
+
+	// Move the last voter without asking for the lease to move.
+	err := tc.Servers[0].DB().AdminRelocateRange(
+		context.Background(),
+		k.AsRawKey(),
+		tc.Targets(3, 1, 0),
+		nil,   /* nonVoterTargets */
+		false, /* transferLeaseToFirstVoter */
+	)
+	require.NoError(t, err)
+	leaseholder, err := tc.FindRangeLeaseHolder(tc.LookupRangeOrFatal(t, k.AsRawKey()), nil /* hint */)
+	require.NoError(t, err)
+	require.Equal(
+		t,
+		roachpb.ReplicationTarget{NodeID: leaseholder.NodeID, StoreID: leaseholder.StoreID},
+		tc.Target(0),
+	)
+}
+
 func TestAdminRelocateRangeFailsWithDuplicates(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -339,7 +385,11 @@ func TestAdminRelocateRangeFailsWithDuplicates(t *testing.T) {
 	}
 	for _, subtest := range tests {
 		err := tc.Servers[0].DB().AdminRelocateRange(
-			context.Background(), k.AsRawKey(), tc.Targets(subtest.voterTargets...), tc.Targets(subtest.nonVoterTargets...),
+			context.Background(),
+			k.AsRawKey(),
+			tc.Targets(subtest.voterTargets...),
+			tc.Targets(subtest.nonVoterTargets...),
+			true, /* transferLeaseToFirstVoter */
 		)
 		require.Regexp(t, subtest.expectedErr, err)
 	}

pkg/kv/kvserver/client_replica_test.go

@@ -2461,9 +2461,14 @@ func TestRandomConcurrentAdminChangeReplicasRequests(t *testing.T) {
 	var wg sync.WaitGroup
 	key := roachpb.Key("a")
 	db := tc.Servers[0].DB()
-	require.Nil(t, db.AdminRelocateRange(
-		ctx, key, makeReplicationTargets(1, 2, 3), nil,
-	))
+	require.Nil(
+		t,
+		db.AdminRelocateRange(
+			ctx, key, makeReplicationTargets(1, 2, 3),
+			nil,  /* nonVoterTargets */
+			true, /* transferLeaseToFirstVoter */
+		),
+	)
 	// Random targets consisting of a random number of nodes from the set of nodes
 	// in the cluster which currently do not have a replica.
 	pickTargets := func() []roachpb.ReplicationTarget {
@@ -3015,8 +3020,13 @@ func TestAdminRelocateRangeSafety(t *testing.T) {
 
 	key := roachpb.Key("a")
 	assert.Nil(t, db.AdminRelocateRange(
-		ctx, key, makeReplicationTargets(1, 2, 3), makeReplicationTargets(),
-	))
+		ctx,
+		key,
+		makeReplicationTargets(1, 2, 3),
+		makeReplicationTargets(),
+		true, /* transferLeaseToFirstVoter */
+	),
+	)
 	rangeInfo, err := getRangeInfo(ctx, db, key)
 	assert.Nil(t, err)
 	assert.Len(t, rangeInfo.Desc.InternalReplicas, 3)
@@ -3052,7 +3062,11 @@ func TestAdminRelocateRangeSafety(t *testing.T) {
 	}
 	relocate := func() {
 		relocateErr = db.AdminRelocateRange(
-			ctx, key, makeReplicationTargets(1, 2, 4), makeReplicationTargets(),
+			ctx,
+			key,
+			makeReplicationTargets(1, 2, 4),
+			makeReplicationTargets(),
+			true, /* transferLeaseToFirstVoter */
 		)
 	}
 	useSeenAdd.Store(true)

pkg/kv/kvserver/merge_queue.go

@@ -349,7 +349,11 @@ func (mq *mergeQueue) process(
 		// The merge queue will only merge ranges that have the same zone config
 		// (see check inside mergeQueue.shouldQueue).
 		if err := mq.store.DB().AdminRelocateRange(
-			ctx, rhsDesc.StartKey, voterTargets, nonVoterTargets,
+			ctx,
+			rhsDesc.StartKey,
+			voterTargets,
+			nonVoterTargets,
+			true, /* transferLeaseToFirstVoter */
 		); err != nil {
 			return false, err
 		}

pkg/kv/kvserver/replica_command.go

@@ -2657,6 +2657,7 @@ func (s *Store) AdminRelocateRange(
 	ctx context.Context,
 	rangeDesc roachpb.RangeDescriptor,
 	voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	transferLeaseToFirstVoter bool,
 ) error {
 	if containsDuplicates(voterTargets) {
 		return errors.AssertionFailedf(
@@ -2686,7 +2687,9 @@ func (s *Store) AdminRelocateRange(
 	}
 	rangeDesc = *newDesc
 
-	rangeDesc, err = s.relocateReplicas(ctx, rangeDesc, voterTargets, nonVoterTargets)
+	rangeDesc, err = s.relocateReplicas(
+		ctx, rangeDesc, voterTargets, nonVoterTargets, transferLeaseToFirstVoter,
+	)
 	if err != nil {
 		return err
 	}
@@ -2718,6 +2721,7 @@ func (s *Store) relocateReplicas(
 	ctx context.Context,
 	rangeDesc roachpb.RangeDescriptor,
 	voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	transferLeaseToFirstVoter bool,
 ) (roachpb.RangeDescriptor, error) {
 	startKey := rangeDesc.StartKey.AsRawKey()
 	transferLease := func(target roachpb.ReplicationTarget) error {
@@ -2744,10 +2748,13 @@ func (s *Store) relocateReplicas(
 				return rangeDesc, err
 			}
 
-			ops, leaseTarget, err := s.relocateOne(ctx, &rangeDesc, voterTargets, nonVoterTargets)
+			ops, leaseTarget, err := s.relocateOne(
+				ctx, &rangeDesc, voterTargets, nonVoterTargets, transferLeaseToFirstVoter,
+			)
 			if err != nil {
 				return rangeDesc, err
 			}
+
 			if leaseTarget != nil {
 				// NB: we may need to transfer even if there are no ops, to make
 				// sure the attempt is made to make the first target the final
@@ -2756,6 +2763,7 @@ func (s *Store) relocateReplicas(
 					return rangeDesc, err
 				}
 			}
+
 			if len(ops) == 0 {
 				// Done.
 				return rangeDesc, ctx.Err()
@@ -2833,6 +2841,7 @@ func (s *Store) relocateOne(
 	ctx context.Context,
 	desc *roachpb.RangeDescriptor,
 	voterTargets, nonVoterTargets []roachpb.ReplicationTarget,
+	transferLeaseToFirstVoter bool,
 ) ([]roachpb.ReplicationChange, *roachpb.ReplicationTarget, error) {
 	if repls := desc.Replicas(); len(repls.VoterFullAndNonVoterDescriptors()) != len(repls.Descriptors()) {
 		// The caller removed all the learners and left the joint config, so there
@@ -3065,9 +3074,9 @@ func (s *Store) relocateOne(
 		}
 	}
 
-	if len(ops) == 0 {
-		// Make sure that the first target is the final leaseholder, as
-		// AdminRelocateRange specifies.
+	if len(ops) == 0 && transferLeaseToFirstVoter {
+		// Make sure that the first target is the final leaseholder, if the caller
+		// asked for it.
 		transferTarget = &voterTargets[0]
 	}
 	return ops, transferTarget, nil

pkg/kv/kvserver/replica_learner_test.go

@@ -976,7 +976,11 @@ func TestLearnerOrJointConfigAdminRelocateRange(t *testing.T) {
 
 	check := func(voterTargets []roachpb.ReplicationTarget) {
 		require.NoError(t, tc.Server(0).DB().AdminRelocateRange(
-			ctx, scratchStartKey, voterTargets, []roachpb.ReplicationTarget{},
+			ctx,
+			scratchStartKey,
+			voterTargets,
+			[]roachpb.ReplicationTarget{},
+			true, /* transferLeaseToFirstVoter */
 		))
 		desc := tc.LookupRangeOrFatal(t, scratchStartKey)
 		voters := desc.Replicas().VoterDescriptors()

pkg/kv/kvserver/replica_send.go

@@ -897,7 +897,16 @@ func (r *Replica) executeAdminBatch(
 		}
 
 	case *roachpb.AdminRelocateRangeRequest:
-		err := r.store.AdminRelocateRange(ctx, *r.Desc(), tArgs.VoterTargets, tArgs.NonVoterTargets)
+		// Transferring the lease to the first voting replica in the target slice is
+		// pre-22.1 behavior.
+		// We revert to that behavior if the request is coming
+		// from a 21.2 node that doesn't yet know about this change in contract.
+		transferLeaseToFirstVoter := !tArgs.TransferLeaseToFirstVoterAccurate
+		// We also revert to that behavior if the caller specifically asked for it.
+		transferLeaseToFirstVoter = transferLeaseToFirstVoter || tArgs.TransferLeaseToFirstVoter
+		err := r.store.AdminRelocateRange(
+			ctx, *r.Desc(), tArgs.VoterTargets, tArgs.NonVoterTargets, transferLeaseToFirstVoter,
+		)
 		pErr = roachpb.NewError(err)
 		resp = &roachpb.AdminRelocateRangeResponse{}
 

pkg/kv/kvserver/store_rebalancer.go

@@ -347,7 +347,13 @@ func (sr *StoreRebalancer) rebalanceStore(
 
 		timeout := sr.rq.processTimeoutFunc(sr.st, replWithStats.repl)
 		if err := contextutil.RunWithTimeout(ctx, "relocate range", timeout, func(ctx context.Context) error {
-			return sr.rq.store.AdminRelocateRange(ctx, *descBeforeRebalance, voterTargets, nonVoterTargets)
+			return sr.rq.store.DB().AdminRelocateRange(
+				ctx,
+				descBeforeRebalance.StartKey.AsRawKey(),
+				voterTargets,
+				nonVoterTargets,
+				true, /* transferLeaseToFirstVoter */
+			)
 		}); err != nil {
 			log.Errorf(ctx, "unable to relocate range to %v: %+v", voterTargets, err)
 			continue

pkg/roachpb/api.proto

@@ -870,6 +870,21 @@ message AdminRelocateRangeRequest {
   repeated ReplicationTarget voter_targets = 2 [(gogoproto.nullable) = false];
   repeated ReplicationTarget non_voter_targets = 3 [(gogoproto.nullable) = false];
   // TODO(a-robinson): Add "reason"/"details" string fields?
+
+  // As of 22.1 (specifically #74077), leaseholder replicas can remove
+  // themselves from the range. This means that now, in a joint state, the
+  // leaseholder that is removing itself chooses the best target replica to
+  // transfer the lease to, all inside of AdminChangeReplicas.
+  //
+  // This means that the pre-22.1 contract of `AdminRelocateRange` to transfer
+  // the lease to the first voter replica isn't required anymore. Only callers
+  // that rely on this contract should set this attribute.
+  bool transfer_lease_to_first_voter = 4;
+  // TODO(aayush): Migration path:
+  // 22.1: Send and consult the attribute.
+  // 22.2: Send but don't consult the attribute.
+  // 23.1: Stop sending or consulting the attribute. Remove this field.
+  bool transfer_lease_to_first_voter_accurate = 5;
 }
 
 message AdminRelocateRangeResponse {

pkg/sql/relocate.go

@@ -128,13 +128,21 @@ func (n *relocateNode) Next(params runParams) (bool, error) {
 		}
 	case tree.RelocateNonVoters:
 		if err := params.p.ExecCfg().DB.AdminRelocateRange(
-			params.ctx, rowKey, existingVoters, relocationTargets,
+			params.ctx,
+			rowKey,
+			existingVoters,
+			relocationTargets,
+			true, /* transferLeaseToFirstVoter */
 		); err != nil {
 			return false, err
 		}
 	case tree.RelocateVoters:
 		if err := params.p.ExecCfg().DB.AdminRelocateRange(
-			params.ctx, rowKey, relocationTargets, existingNonVoters,
+			params.ctx,
+			rowKey,
+			relocationTargets,
+			existingNonVoters,
+			true, /* transferLeaseToFirstVoter */
 		); err != nil {
 			return false, err
 		}