kvserver: add cross-region snapshot byte metrics to StoreMetrics

pkg/kv/kvserver/allocator/storepool/store_pool.go

@@ -1349,16 +1349,43 @@ func (sp *StorePool) GetLocalitiesPerReplica(
 	return localities
 }
 
+// getNodeLocalityWithString returns the locality information and the string
+// format for the given node.
+func (sp *StorePool) getNodeLocalityWithString(nodeID roachpb.NodeID) localityWithString {
+	nodeLocality := localityWithString{}
+	sp.localitiesMu.RLock()
+	defer sp.localitiesMu.RUnlock()
+	if locality, ok := sp.localitiesMu.nodeLocalities[nodeID]; ok {
+		nodeLocality = locality
+	}
+	// Return an empty localityWithString struct if nothing is found.
+	return nodeLocality
+}
+
 // GetNodeLocalityString returns the locality information for the given node
 // in its string format.
 func (sp *StorePool) GetNodeLocalityString(nodeID roachpb.NodeID) string {
-	sp.localitiesMu.RLock()
-	defer sp.localitiesMu.RUnlock()
-	locality, ok := sp.localitiesMu.nodeLocalities[nodeID]
-	if !ok {
-		return ""
+	return sp.getNodeLocalityWithString(nodeID).str
+}
+
+// getNodeLocality returns the locality information for the given node.
+func (sp *StorePool) getNodeLocality(nodeID roachpb.NodeID) roachpb.Locality {
+	return sp.getNodeLocalityWithString(nodeID).locality
+}
+
+// IsCrossRegion takes in two replicas and compares the locality of them based
+// on their replica node IDs. It returns (bool, error) indicating whether the
+// two replicas’ nodes are in different regions and if any errors occurred
+// during the lookup process.
+func (sp *StorePool) IsCrossRegion(
+	firstReplica roachpb.ReplicaDescriptor, secReplica roachpb.ReplicaDescriptor,
+) (bool, error) {
+	isCrossRegion, err := sp.getNodeLocality(firstReplica.NodeID).IsCrossRegion(
+		sp.getNodeLocality(secReplica.NodeID))
+	if err != nil {
+		return false, err
 	}
-	return locality.str
+	return isCrossRegion, nil
 }
 
 // IsStoreReadyForRoutineReplicaTransfer returns true iff the store's node is

pkg/kv/kvserver/metrics.go

@@ -954,6 +954,18 @@ var (
 		Measurement: "Snapshots",
 		Unit:        metric.Unit_COUNT,
 	}
+	metaRangeSnapShotCrossRegionSentBytes = metric.Metadata{
+		Name:        "range.snapshots.cross-region.sent-bytes",
+		Help:        "Number of snapshot bytes sent cross region",
+		Measurement: "Bytes",
+		Unit:        metric.Unit_BYTES,
+	}
+	metaRangeSnapShotCrossRegionRcvdBytes = metric.Metadata{
+		Name:        "range.snapshots.cross-region.rcvd-bytes",
+		Help:        "Number of snapshot bytes received cross region",
+		Measurement: "Bytes",
+		Unit:        metric.Unit_BYTES,
+	}
 	metaRangeSnapshotSendQueueLength = metric.Metadata{
 		Name:        "range.snapshots.send-queue",
 		Help:        "Number of snapshots queued to send",
@@ -2181,6 +2193,8 @@ type StoreMetrics struct {
 	RangeSnapshotRebalancingSentBytes            *metric.Counter
 	RangeSnapshotRecvFailed                      *metric.Counter
 	RangeSnapshotRecvUnusable                    *metric.Counter
+	RangeSnapShotCrossRegionSentBytes            *metric.Counter
+	RangeSnapShotCrossRegionRcvdBytes            *metric.Counter
 
 	// Range snapshot queue metrics.
 	RangeSnapshotSendQueueLength     *metric.Gauge
@@ -2803,6 +2817,8 @@ func newStoreMetrics(histogramWindow time.Duration) *StoreMetrics {
 		RangeSnapshotRebalancingSentBytes:            metric.NewCounter(metaRangeSnapshotRebalancingSentBytes),
 		RangeSnapshotRecvFailed:                      metric.NewCounter(metaRangeSnapshotRecvFailed),
 		RangeSnapshotRecvUnusable:                    metric.NewCounter(metaRangeSnapshotRecvUnusable),
+		RangeSnapShotCrossRegionSentBytes:            metric.NewCounter(metaRangeSnapShotCrossRegionSentBytes),
+		RangeSnapShotCrossRegionRcvdBytes:            metric.NewCounter(metaRangeSnapShotCrossRegionRcvdBytes),
 		RangeSnapshotSendQueueLength:                 metric.NewGauge(metaRangeSnapshotSendQueueLength),
 		RangeSnapshotRecvQueueLength:                 metric.NewGauge(metaRangeSnapshotRecvQueueLength),
 		RangeSnapshotSendInProgress:                  metric.NewGauge(metaRangeSnapshotSendInProgress),

pkg/kv/kvserver/replica_command.go

@@ -3151,6 +3151,9 @@ func (r *Replica) followerSendSnapshot(
 			r.store.metrics.DelegateSnapshotSendBytes.Inc(inc)
 		}
 		r.store.metrics.RangeSnapshotSentBytes.Inc(inc)
+		if r.store.shouldIncrementCrossRegionSnapshotMetrics(ctx, req.CoordinatorReplica, req.RecipientReplica) {
+			r.store.metrics.RangeSnapShotCrossRegionSentBytes.Inc(inc)
+		}
 
 		switch header.Priority {
 		case kvserverpb.SnapshotRequest_RECOVERY: