server: add StoreIDs to SpanStats response

This commit adds a new response field to the
SpanStats rpc, `StoreIDs` which is a list of
unique store ids in ascending order that store
the requestd span.

Fixes: #129060
Release note: None

pkg/roachpb/span_stats.proto

@@ -70,7 +70,11 @@ message SpanStats {
   // stale.
   cockroach.storage.enginepb.MVCCStats approximate_total_stats = 7 [(gogoproto.nullable) = false];
 
-  // NEXT ID: 8.
+
+  // Unique store ids for the requested span.
+  repeated int32 store_ids = 8 [(gogoproto.customname) = "StoreIDs", (gogoproto.casttype) = "StoreID"];
+
+  // NEXT ID: 9.
 }
 
 message SpanStatsResponse {

pkg/server/span_stats_server.go

@@ -13,6 +13,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"sort"
 	"strconv"
 	"time"
 
@@ -25,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/rangedesc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -151,6 +153,7 @@ func (s *systemStatusServer) spanStatsFanOut(
 			res.SpanToStats[spanStr].ApproximateDiskBytes += spanStats.ApproximateDiskBytes
 			res.SpanToStats[spanStr].RemoteFileBytes += spanStats.RemoteFileBytes
 			res.SpanToStats[spanStr].ExternalFileBytes += spanStats.ExternalFileBytes
+			res.SpanToStats[spanStr].StoreIDs = util.CombineUnique(res.SpanToStats[spanStr].StoreIDs, spanStats.StoreIDs)
 
 			// Logical values: take the values from the node that responded first.
 			// TODO: This should really be read from the leaseholder.
@@ -214,6 +217,7 @@ func (s *systemStatusServer) statsForSpan(
 
 	// First, get the approximate disk bytes from each store.
 	err = s.stores.VisitStores(func(store *kvserver.Store) error {
+		spanStats.StoreIDs = append(spanStats.StoreIDs, store.StoreID())
 		approxDiskBytes, remoteBytes, externalBytes, err := store.TODOEngine().ApproximateDiskBytes(rSpan.Key.AsRawKey(), rSpan.EndKey.AsRawKey())
 		if err != nil {
 			return err
@@ -224,6 +228,10 @@ func (s *systemStatusServer) statsForSpan(
 		return nil
 	})
 
+	sort.Slice(spanStats.StoreIDs, func(i, j int) bool {
+		return spanStats.StoreIDs[i] < spanStats.StoreIDs[j]
+	})
+
 	if err != nil {
 		return nil, err
 	}
@@ -287,6 +295,7 @@ func (s *systemStatusServer) statsForSpan(
 			log.VEventf(ctx, 1, "Range %v exceeds span %v, calculating stats for subspan %v",
 				descSpan, rSpan, roachpb.RSpan{Key: scanStart, EndKey: scanEnd},
 			)
+
 			err = s.stores.VisitStores(func(s *kvserver.Store) error {
 				stats, err := storage.ComputeStats(
 					ctx,
@@ -303,7 +312,6 @@ func (s *systemStatusServer) statsForSpan(
 				spanStats.TotalStats.Add(stats)
 				return nil
 			})
-
 			if err != nil {
 				return nil, err
 			}

pkg/server/span_stats_test.go

@@ -162,13 +162,33 @@ func TestSpanStatsFanOut(t *testing.T) {
 		)
 		require.NoError(t, err)
 
+		equalStoreIDs := func(a, b []roachpb.StoreID) bool {
+			if len(a) != len(b) {
+				return false
+			}
+			for i := range a {
+				if a[i] != b[i] {
+					return false
+				}
+			}
+			return true
+		}
+
 		// Verify stats across different spans.
 		for _, tcase := range testCases {
 			rSpan, err := keys.SpanAddr(tcase.span)
 			require.NoError(t, err)
 
 			// Assert expected values from multi-span request
 			spanStats := multiResult.SpanToStats[tcase.span.String()]
+			if !equalStoreIDs([]roachpb.StoreID{1, 2, 3}, spanStats.StoreIDs) {
+				return errors.Newf("Multi-span: expected storeIDs %v in span [%s - %s], found %v",
+					[]roachpb.StoreID{1},
+					rSpan.Key.String(),
+					rSpan.EndKey.String(),
+					spanStats.StoreIDs,
+				)
+			}
 			if tcase.expectedRanges != spanStats.RangeCount {
 				return errors.Newf("Multi-span: expected %d ranges in span [%s - %s], found %d",
 					tcase.expectedRanges,