From 3b99b1f071135f077af1518e4bea37983dc68e14 Mon Sep 17 00:00:00 2001
From: Andrei Matei <andrei@cockroachlabs.com>
Date: Wed, 9 Sep 2020 18:51:57 -0400
Subject: [PATCH] plumb ctx

---
 pkg/kv/kvserver/replica_command.go | 20 ++++++++++++--------
 pkg/kv/kvserver/replica_test.go    |  2 +-
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/pkg/kv/kvserver/replica_command.go b/pkg/kv/kvserver/replica_command.go
index 52be6a73a65f..c35898335f4a 100644
--- a/pkg/kv/kvserver/replica_command.go
+++ b/pkg/kv/kvserver/replica_command.go
@@ -190,7 +190,7 @@ func splitTxnAttempt(
 	{
 		b := txn.NewBatch()
 		leftDescKey := keys.RangeDescriptorKey(leftDesc.StartKey)
-		if err := updateRangeDescriptor(b, leftDescKey, dbDescValue, leftDesc); err != nil {
+		if err := updateRangeDescriptor(ctx, b, leftDescKey, dbDescValue, leftDesc); err != nil {
 			return err
 		}
 		// Commit this batch first to ensure that the transaction record
@@ -215,7 +215,7 @@ func splitTxnAttempt(
 	// Write range descriptor for right hand side of the split.
 	rightDescKey := keys.RangeDescriptorKey(rightDesc.StartKey)
 	log.Infof(ctx, "!!! will write RHS desc: %s", rightDesc)
-	if err := updateRangeDescriptor(b, rightDescKey, nil, rightDesc); err != nil {
+	if err := updateRangeDescriptor(ctx, b, rightDescKey, nil, rightDesc); err != nil {
 		return err
 	}
 
@@ -253,7 +253,7 @@ func splitTxnStickyUpdateAttempt(
 
 	b := txn.NewBatch()
 	descKey := keys.RangeDescriptorKey(desc.StartKey)
-	if err := updateRangeDescriptor(b, descKey, dbDescValue, &newDesc); err != nil {
+	if err := updateRangeDescriptor(ctx, b, descKey, dbDescValue, &newDesc); err != nil {
 		return err
 	}
 	if err := updateRangeAddressing(b, &newDesc); err != nil {
@@ -465,7 +465,7 @@ func (r *Replica) adminUnsplitWithDescriptor(
 		descKey := keys.RangeDescriptorKey(newDesc.StartKey)
 
 		b := txn.NewBatch()
-		if err := updateRangeDescriptor(b, descKey, dbDescValue, &newDesc); err != nil {
+		if err := updateRangeDescriptor(ctx, b, descKey, dbDescValue, &newDesc); err != nil {
 			return err
 		}
 		if err := updateRangeAddressing(b, &newDesc); err != nil {
@@ -672,7 +672,7 @@ func (r *Replica) AdminMerge(
 			b := txn.NewBatch()
 			leftDescKey := keys.RangeDescriptorKey(updatedLeftDesc.StartKey)
 			if err := updateRangeDescriptor(
-				b, leftDescKey, dbOrigLeftDescValue, &updatedLeftDesc,
+				ctx, b, leftDescKey, dbOrigLeftDescValue, &updatedLeftDesc,
 			); err != nil {
 				return err
 			}
@@ -701,7 +701,7 @@ func (r *Replica) AdminMerge(
 		}
 
 		// Remove the range descriptor for the deleted range.
-		if err := updateRangeDescriptor(b, rightDescKey, dbRightDescKV.Value, nil); err != nil {
+		if err := updateRangeDescriptor(ctx, b, rightDescKey, dbRightDescKV.Value, nil); err != nil {
 			return err
 		}
 
@@ -1608,7 +1608,7 @@ func execChangeReplicasTxn(
 
 			// Important: the range descriptor must be the first thing touched in the transaction
 			// so the transaction record is co-located with the range being modified.
-			if err := updateRangeDescriptor(b, descKey, dbDescValue, crt.Desc); err != nil {
+			if err := updateRangeDescriptor(ctx, b, descKey, dbDescValue, crt.Desc); err != nil {
 				return err
 			}
 
@@ -2007,7 +2007,11 @@ func conditionalGetDescValueFromDB(
 // descriptor, a CommitTrigger must be used to update the in-memory
 // descriptor; it will not automatically be copied from newDesc.
 func updateRangeDescriptor(
-	b *kv.Batch, descKey roachpb.Key, oldValue *roachpb.Value, newDesc *roachpb.RangeDescriptor,
+	ctx context.Context,
+	b *kv.Batch,
+	descKey roachpb.Key,
+	oldValue *roachpb.Value,
+	newDesc *roachpb.RangeDescriptor,
 ) error {
 	// This is subtle: []byte(nil) != interface{}(nil). A []byte(nil) refers to
 	// an empty value. An interface{}(nil) refers to a non-existent value. So
diff --git a/pkg/kv/kvserver/replica_test.go b/pkg/kv/kvserver/replica_test.go
index d177de7f15f0..f277082d46b2 100644
--- a/pkg/kv/kvserver/replica_test.go
+++ b/pkg/kv/kvserver/replica_test.go
@@ -398,7 +398,7 @@ func (tc *testContext) addBogusReplicaToRangeDesc(
 		Header: roachpb.Header{Timestamp: tc.Clock().Now()},
 	}
 	descKey := keys.RangeDescriptorKey(oldDesc.StartKey)
-	if err := updateRangeDescriptor(&ba, descKey, dbDescKV.Value, &newDesc); err != nil {
+	if err := updateRangeDescriptor(ctx, &ba, descKey, dbDescKV.Value, &newDesc); err != nil {
 		return roachpb.ReplicaDescriptor{}, err
 	}
 	if err := tc.store.DB().Run(ctx, &ba); err != nil {