From c05165486d61eda81268d4681c6e6661368db2ab Mon Sep 17 00:00:00 2001
From: Wenyi <wenyi.hu@cockroachlabs.com>
Date: Fri, 30 Jun 2023 07:58:58 -0400
Subject: [PATCH] to be squashed

---
 pkg/kv/kvserver/replica_command.go | 4 ++--
 pkg/kv/kvserver/store_snapshot.go  | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/pkg/kv/kvserver/replica_command.go b/pkg/kv/kvserver/replica_command.go
index b9fb52fc1b4e..06b90bdcf7c3 100644
--- a/pkg/kv/kvserver/replica_command.go
+++ b/pkg/kv/kvserver/replica_command.go
@@ -3151,6 +3151,8 @@ func (r *Replica) followerSendSnapshot(
 	sent := func() {
 		r.store.metrics.RangeSnapshotsGenerated.Inc(1)
 	}
+	comparisonResult := r.store.getLocalityComparison(ctx, req.CoordinatorReplica.NodeID,
+		req.RecipientReplica.NodeID)
 
 	recordBytesSent := func(inc int64) {
 		// Only counts for delegated bytes if we are not self-delegating.
@@ -3158,8 +3160,6 @@ func (r *Replica) followerSendSnapshot(
 			r.store.metrics.DelegateSnapshotSendBytes.Inc(inc)
 		}
 		r.store.metrics.RangeSnapshotSentBytes.Inc(inc)
-		comparisonResult := r.store.getLocalityComparison(ctx, req.CoordinatorReplica.NodeID,
-			req.RecipientReplica.NodeID)
 		r.store.metrics.updateCrossLocalityMetricsOnSnapshotSent(comparisonResult, inc)
 
 		switch header.Priority {
diff --git a/pkg/kv/kvserver/store_snapshot.go b/pkg/kv/kvserver/store_snapshot.go
index a36060e1bbb1..ee1ea28e198c 100644
--- a/pkg/kv/kvserver/store_snapshot.go
+++ b/pkg/kv/kvserver/store_snapshot.go
@@ -1105,10 +1105,11 @@ func (s *Store) receiveSnapshot(
 		log.Infof(ctx, "accepted snapshot reservation for r%d", header.State.Desc.RangeID)
 	}
 
+	comparisonResult := s.getLocalityComparison(ctx,
+		header.RaftMessageRequest.FromReplica.NodeID, header.RaftMessageRequest.ToReplica.NodeID)
+
 	recordBytesReceived := func(inc int64) {
 		s.metrics.RangeSnapshotRcvdBytes.Inc(inc)
-		comparisonResult := s.getLocalityComparison(ctx,
-			header.RaftMessageRequest.FromReplica.NodeID, header.RaftMessageRequest.ToReplica.NodeID)
 		s.metrics.updateCrossLocalityMetricsOnSnapshotRcvd(comparisonResult, inc)
 
 		switch header.Priority {