From de3b318111099d397b6822cb24ad82ae2ecb47bf Mon Sep 17 00:00:00 2001
From: Thomas Hardy <thardy@cockroachlabs.com>
Date: Wed, 2 Aug 2023 13:25:31 -0400
Subject: [PATCH] server: return authoritative span statistics for db details
 endpoint

Resolves: #96163

This change makes the admin API endpoint getting database statistics
scan KV for span statistics instead of using the range descriptor cache.
This provides authoritative output, helping deflake
`TestMultiRegionDatabaseStats`.

Release note (sql change): admin API database details endpoint now
returns authoritative range statistics.
---
 pkg/ccl/serverccl/admin_test.go |  6 +++++
 pkg/server/admin.go             | 39 ++++++++++++++++-----------------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/pkg/ccl/serverccl/admin_test.go b/pkg/ccl/serverccl/admin_test.go
index 07a90515a766..958db1707605 100644
--- a/pkg/ccl/serverccl/admin_test.go
+++ b/pkg/ccl/serverccl/admin_test.go
@@ -225,6 +225,12 @@ func TestTableAndDatabaseDetailsAndStats(t *testing.T) {
 
 	require.Equal(t, dbResp.TableNames[0], "public.test")
 
+	var dbDetailsResp serverpb.DatabaseDetailsResponse
+	err = getAdminJSONProto(st, "databases/defaultdb?include_stats=true", &dbDetailsResp)
+	require.NoError(t, err)
+
+	require.Greater(t, dbDetailsResp.Stats.RangeCount, int64(0))
+
 	// TableStats
 	tableStatsResp := &serverpb.TableStatsResponse{}
 	err = getAdminJSONProto(st, "databases/defaultdb/tables/public.test/stats", tableStatsResp)
diff --git a/pkg/server/admin.go b/pkg/server/admin.go
index ba0d7035df6b..543dfcf58f4a 100644
--- a/pkg/server/admin.go
+++ b/pkg/server/admin.go
@@ -67,6 +67,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
+	"github.com/cockroachdb/cockroach/pkg/util/rangedesc"
 	"github.com/cockroachdb/cockroach/pkg/util/safesql"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -713,6 +714,7 @@ func (s *adminServer) getDatabaseStats(
 		select {
 		case response := <-responses:
 			if response.err != nil {
+				fmt.Printf("database stats - encountered error getting table span stats: %+v\n", response.err)
 				stats.MissingTables = append(
 					stats.MissingTables,
 					&serverpb.DatabaseDetailsResponse_Stats_MissingTable{
@@ -1328,16 +1330,9 @@ func (s *adminServer) statsForSpan(
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	rSpan, err := keys.SpanAddr(span)
-	if err != nil {
-		return nil, err
-	}
-
 	// Get a list of nodeIDs, range counts, and replica counts per node
 	// for the specified span.
-	nodeIDs, rangeCount, replCounts, err := getNodeIDsRangeCountReplCountForSpan(
-		ctx, s.distSender, rSpan,
-	)
+	nodeIDs, rangeCount, replCounts, err := s.getSpanDetails(ctx, span)
 	if err != nil {
 		return nil, err
 	}
@@ -1450,25 +1445,29 @@ func (s *adminServer) statsForSpan(
 
 // Returns the list of node ids, range count,
 // and replica count for the specified span.
-func getNodeIDsRangeCountReplCountForSpan(
-	ctx context.Context, ds *kvcoord.DistSender, rSpan roachpb.RSpan,
+func (s *adminServer) getSpanDetails(
+	ctx context.Context, span roachpb.Span,
 ) (nodeIDList []roachpb.NodeID, rangeCount int64, replCounts map[roachpb.NodeID]int64, _ error) {
 	nodeIDs := make(map[roachpb.NodeID]struct{})
 	replCountForNodeID := make(map[roachpb.NodeID]int64)
-	ri := kvcoord.MakeRangeIterator(ds)
-	ri.Seek(ctx, rSpan.Key, kvcoord.Ascending)
-	for ; ri.Valid(); ri.Next(ctx) {
+	var it rangedesc.Iterator
+	var err error
+	if s.sqlServer.tenantConnect == nil {
+		it, err = s.sqlServer.execCfg.RangeDescIteratorFactory.NewIterator(ctx, span)
+	} else {
+		it, err = s.sqlServer.tenantConnect.NewIterator(ctx, span)
+	}
+	if err != nil {
+		return nil, 0, nil, err
+	}
+	var rangeDesc roachpb.RangeDescriptor
+	for ; it.Valid(); it.Next() {
 		rangeCount++
-		for _, repl := range ri.Desc().Replicas().Descriptors() {
+		rangeDesc = it.CurRangeDescriptor()
+		for _, repl := range rangeDesc.Replicas().Descriptors() {
 			replCountForNodeID[repl.NodeID]++
 			nodeIDs[repl.NodeID] = struct{}{}
 		}
-		if !ri.NeedAnother(rSpan) {
-			break
-		}
-	}
-	if err := ri.Error(); err != nil {
-		return nil, 0, nil, err
 	}
 
 	nodeIDList = make([]roachpb.NodeID, 0, len(nodeIDs))