kvclient, server: implement SpanStats for use with coalesced ranges

This commit provides a re-implementation of fetching the MVCC stats and range count for a given span. The previous implementation relied on the assumption that 1 range can contain at most 1 table/index. Since cockroachdbGH-79700, this assumption is no longer true. This re-implementation has three components. The first is spanstats.Accessor, which provides an interface for accessing span stats from KV. The second is spanstatsaccessor.LocalAccessor, which provides an implementation of the interface for callers that are co-located on a KV node. The third is an implementation by kvtenantccl.Connector, which provides an implementation of the interface for non co-located callers, like SQL pods. The interaction between these components is illustrated below: System Tenant +----------------------------------------------------+ | | | | | KV Node fan-out | | +----------+ | | | | | | +-------------------------------v----+ | | | | | | | | | serverpb.InternalSpanStatsServer +-----+ | | | | | | +----------------+-------------------+ | | | | | roachpb.InternalSpanStatsResponse | | | | | | | | v | | +----------------------------------+ | | | | | | | spanstatsaccessor.LocalAccessor | | | | | | | +-----------------+----------------+ | | | | | +---------------------+ | | roachpb.InternalSpanStatsResponse | | | | | | | v v | | +-------------------------+ +-----------+ | | | | | | | | | serverpb.StatusServer | | SQLServer | | | | | | | | | +------------+------------+ +-----------+ | | | | | | | +----------------------+-----------------------------+ | | serverpb.SpanStatsResponse | | | | | Secondary Tenant | +----------------------+------------------------------+ | | | | +------------v-------------+ | | | | | | | kvtenantccl.Connector | | | | | | | +------------+-------------+ | | | | | | | | +--------------------+ | | roachpb.InternalSpanStatsResponse| | | | | | | | | | | +-----------------v-------+ +------v------+ | | | | | | | | | serverpb.StatusServer | | SQLServer | | | | | | | | | +-------------------------+ +-------------+ | | | | | | | +-----------------------------------------------------+ Resolves cockroachdb#84105 Part of: https://cockroachlabs.atlassian.net/browse/CRDB-22711 Release note (backward-incompatible change): The SpanStatsRequest message field 'node_id' has changed from type 'string' to type 'int32'.
zachlite · Feb 7, 2023 · 0516358 · 0516358
1 parent 9444613
commit 0516358
Show file tree

Hide file tree

Showing 31 changed files with 706 additions and 240 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -391,6 +391,7 @@
 /pkg/keyvisualizer/          @cockroachdb/kv-obs-prs
 /pkg/multitenant/            @cockroachdb/multi-tenant
 /pkg/release/                @cockroachdb/dev-inf
+/pkg/roachpb/                @cockroachdb/kv-prs
 /pkg/roachpb/.gitattributes  @cockroachdb/dev-inf
 /pkg/roachpb/ambiguous_*     @cockroachdb/kv-prs
 /pkg/roachpb/api*            @cockroachdb/kv-prs

diff --git a/docs/generated/http/full.md b/docs/generated/http/full.md
@@ -2936,7 +2936,7 @@ Support status: [reserved](#support-status)
 
 | Field | Type | Label | Description | Support status |
 | ----- | ---- | ----- | ----------- | -------------- |
-| node_id | [string](#cockroach.server.serverpb.SpanStatsRequest-string) |  |  | [reserved](#support-status) |
+| node_id | [int32](#cockroach.server.serverpb.SpanStatsRequest-int32) |  |  | [reserved](#support-status) |
 | start_key | [bytes](#cockroach.server.serverpb.SpanStatsRequest-bytes) |  |  | [reserved](#support-status) |
 | end_key | [bytes](#cockroach.server.serverpb.SpanStatsRequest-bytes) |  |  | [reserved](#support-status) |
 

diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel
@@ -1159,6 +1159,8 @@ GO_TARGETS = [
     "//pkg/kv/kvclient/rangefeed:rangefeed",
     "//pkg/kv/kvclient/rangefeed:rangefeed_test",
     "//pkg/kv/kvclient/rangestats:rangestats",
+    "//pkg/kv/kvclient/spanstats/spanstatsaccessor:spanstatsaccessor",
+    "//pkg/kv/kvclient/spanstats:spanstats",
     "//pkg/kv/kvclient:kvclient",
     "//pkg/kv/kvnemesis/kvnemesisutil:kvnemesisutil",
     "//pkg/kv/kvnemesis:kvnemesis",
@@ -2581,6 +2583,8 @@ GET_X_DATA_TARGETS = [
     "//pkg/kv/kvclient/rangefeed/rangefeedbuffer:get_x_data",
     "//pkg/kv/kvclient/rangefeed/rangefeedcache:get_x_data",
     "//pkg/kv/kvclient/rangestats:get_x_data",
+    "//pkg/kv/kvclient/spanstats:get_x_data",
+    "//pkg/kv/kvclient/spanstats/spanstatsaccessor:get_x_data",
     "//pkg/kv/kvnemesis:get_x_data",
     "//pkg/kv/kvnemesis/kvnemesisutil:get_x_data",
     "//pkg/kv/kvprober:get_x_data",

diff --git a/pkg/ccl/kvccl/kvtenantccl/connector.go b/pkg/ccl/kvccl/kvtenantccl/connector.go
@@ -648,6 +648,27 @@ func (c *Connector) SpanConfigConformance(
 	return report, nil
 }
 
+// SpanStats implements the spanstats.Accessor Interface.
+func (c *Connector) SpanStats(
+	ctx context.Context, startKey,
+	endKey roachpb.Key, nodeID roachpb.NodeID,
+) (*roachpb.InternalSpanStatsResponse, error) {
+	var response *roachpb.InternalSpanStatsResponse
+	err := c.withClient(ctx, func(ctx context.Context, c *client) error {
+		stats, err := c.SpanStats(ctx, &serverpb.SpanStatsRequest{
+			NodeID:   nodeID,
+			StartKey: roachpb.RKey(startKey),
+			EndKey:   roachpb.RKey(endKey),
+		})
+		if err != nil {
+			return err
+		}
+		response = (*roachpb.InternalSpanStatsResponse)(stats)
+		return nil
+	})
+	return response, err
+}
+
 // GetAllSystemSpanConfigsThatApply implements the spanconfig.KVAccessor
 // interface.
 func (c *Connector) GetAllSystemSpanConfigsThatApply(

diff --git a/pkg/kv/kvclient/kvtenant/BUILD.bazel b/pkg/kv/kvclient/kvtenant/BUILD.bazel
@@ -12,6 +12,7 @@ go_library(
         "//pkg/keys",
         "//pkg/kv/kvclient/kvcoord",
         "//pkg/kv/kvclient/rangecache",
+        "//pkg/kv/kvclient/spanstats",
         "//pkg/roachpb",
         "//pkg/rpc",
         "//pkg/rpc/nodedialer",

diff --git a/pkg/kv/kvclient/kvtenant/connector.go b/pkg/kv/kvclient/kvtenant/connector.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangecache"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/spanstats"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
@@ -91,6 +92,9 @@ type Connector interface {
 	// applicable to secondary tenants.
 	spanconfig.KVAccessor
 
+	// Accessor provides access to span stats.
+	spanstats.Accessor
+
 	// Reporter provides access to conformance reports, i.e. whether ranges
 	// backing queried keyspans conform the span configs that apply to them.
 	spanconfig.Reporter

diff --git a/pkg/kv/kvclient/spanstats/BUILD.bazel b/pkg/kv/kvclient/spanstats/BUILD.bazel
@@ -0,0 +1,12 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "spanstats",
+    srcs = ["span_stats.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvclient/spanstats",
+    visibility = ["//visibility:public"],
+    deps = ["//pkg/roachpb"],
+)
+
+get_x_data(name = "get_x_data")
diff --git a/pkg/kv/kvclient/spanstats/span_stats.go b/pkg/kv/kvclient/spanstats/span_stats.go
@@ -0,0 +1,32 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package spanstats
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+)
+
+// Accessor mediates access to SpanStats by consumers across the KV/Tenant
+// boundary.
+type Accessor interface {
+	// SpanStats returns the span stats between startKey and endKey.
+	// The caller can request span stats from across the cluster by setting
+	// nodeId to 0. The caller can request span stats from a specific node by
+	// setting the value of nodeID accordingly.
+	SpanStats(
+		ctx context.Context,
+		startKey,
+		endKey roachpb.Key,
+		nodeID roachpb.NodeID,
+	) (*roachpb.InternalSpanStatsResponse, error)
+}
diff --git a/pkg/kv/kvclient/spanstats/spanstatsaccessor/BUILD.bazel b/pkg/kv/kvclient/spanstats/spanstatsaccessor/BUILD.bazel
@@ -0,0 +1,15 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "spanstatsaccessor",
+    srcs = ["accessor.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvclient/spanstats/spanstatsaccessor",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/roachpb",
+        "//pkg/server/serverpb",
+    ],
+)
+
+get_x_data(name = "get_x_data")
diff --git a/pkg/kv/kvclient/spanstats/spanstatsaccessor/accessor.go b/pkg/kv/kvclient/spanstats/spanstatsaccessor/accessor.go
@@ -0,0 +1,51 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package spanstatsaccessor
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+)
+
+// LocalAccessor is an implementation of spanstats.Accessor that is meant
+// to provide access to span stats to servers that are co-located on the KV
+// node.
+type LocalAccessor struct {
+	spanStatsServer serverpb.InternalSpanStatsServer
+}
+
+// New returns a new instance of AccessorImpl.
+func New(server serverpb.InternalSpanStatsServer) *LocalAccessor {
+	return &LocalAccessor{spanStatsServer: server}
+}
+
+// SpanStats implements the spanstats.Accessor interface.
+func (a *LocalAccessor) SpanStats(
+	ctx context.Context, startKey,
+	endKey roachpb.Key, nodeID roachpb.NodeID,
+) (*roachpb.InternalSpanStatsResponse, error) {
+	res, err := a.spanStatsServer.GetSpanStatsInternal(ctx,
+		&roachpb.InternalSpanStatsRequest{
+			Span: roachpb.Span{
+				Key:    startKey,
+				EndKey: endKey,
+			},
+			NodeID: nodeID,
+		})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
diff --git a/pkg/kv/kvserver/BUILD.bazel b/pkg/kv/kvserver/BUILD.bazel
@@ -248,7 +248,6 @@ go_test(
         "client_spanconfigs_test.go",
         "client_split_burst_test.go",
         "client_split_test.go",
-        "client_status_test.go",
         "client_store_test.go",
         "client_tenant_test.go",
         "client_test.go",

diff --git a/pkg/kv/kvserver/store.go b/pkg/kv/kvserver/store.go
@@ -63,7 +63,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/spanconfig"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigstore"
 	"github.com/cockroachdb/cockroach/pkg/storage"
-	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/contextutil"
@@ -3138,33 +3137,6 @@ func mapToHotReplicasInfo(repls []CandidateReplica) []HotReplicaInfo {
 	return hotRepls
 }
 
-// StoreKeySpanStats carries the result of a stats computation over a key range.
-type StoreKeySpanStats struct {
-	ReplicaCount         int
-	MVCC                 enginepb.MVCCStats
-	ApproximateDiskBytes uint64
-}
-
-// ComputeStatsForKeySpan computes the aggregated MVCCStats for all replicas on
-// this store which contain any keys in the supplied range.
-func (s *Store) ComputeStatsForKeySpan(startKey, endKey roachpb.RKey) (StoreKeySpanStats, error) {
-	var result StoreKeySpanStats
-
-	newStoreReplicaVisitor(s).UndefinedOrder().Visit(func(repl *Replica) bool {
-		desc := repl.Desc()
-		if bytes.Compare(startKey, desc.EndKey) >= 0 || bytes.Compare(desc.StartKey, endKey) >= 0 {
-			return true // continue
-		}
-		result.MVCC.Add(repl.GetMVCCStats())
-		result.ReplicaCount++
-		return true
-	})
-
-	var err error
-	result.ApproximateDiskBytes, err = s.engine.ApproximateDiskBytes(startKey.AsRawKey(), endKey.AsRawKey())
-	return result, err
-}
-
 // ReplicateQueueDryRun runs the given replica through the replicate queue
 // (using the allocator) without actually carrying out any changes, returning
 // all trace messages collected along the way.

diff --git a/pkg/roachpb/BUILD.bazel b/pkg/roachpb/BUILD.bazel
@@ -141,6 +141,7 @@ proto_library(
         "io-formats.proto",
         "metadata.proto",
         "span_config.proto",
+        "span_stats.proto",
     ],
     strip_import_prefix = "/pkg",
     visibility = ["//visibility:public"],
@@ -158,6 +159,7 @@ proto_library(
         "@com_github_gogo_protobuf//gogoproto:gogo_proto",
         "@com_google_protobuf//:duration_proto",
         "@com_google_protobuf//:timestamp_proto",
+        "@go_googleapis//google/api:annotations_proto",
     ],
 )
 
@@ -179,6 +181,7 @@ go_proto_library(
         "//pkg/util/tracing/tracingpb",
         "@com_github_cockroachdb_errors//errorspb",
         "@com_github_gogo_protobuf//gogoproto",
+        "@org_golang_google_genproto//googleapis/api/annotations:go_default_library",
     ],
 )
 

diff --git a/pkg/roachpb/go.mod b/pkg/roachpb/go.mod
@@ -0,0 +1,3 @@
+module github.com/cockroachdb/cockroach/pkg/roachpb
+
+go 1.19
diff --git a/pkg/roachpb/span_stats.proto b/pkg/roachpb/span_stats.proto
@@ -0,0 +1,35 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+syntax = "proto3";
+package cockroach.roachpb;
+option go_package = "roachpb";
+
+import "roachpb/data.proto";
+import "storage/enginepb/mvcc.proto";
+import "gogoproto/gogo.proto";
+import "google/api/annotations.proto";
+
+// InternalSpanStatsRequest
+message InternalSpanStatsRequest {
+  Span span = 1 [(gogoproto.nullable) = false];
+
+  // A node_id of `0` indicates the server should issue a fan-out to all nodes.
+  int32 node_id = 2 [(gogoproto.customname) = "NodeID",
+                     (gogoproto.casttype) = "NodeID"];
+}
+
+// InternalSpanStatsResponse
+message InternalSpanStatsResponse {
+  int32 range_count = 2;
+  uint64 approximate_disk_bytes = 3;
+  cockroach.storage.enginepb.MVCCStats total_stats = 1
+      [ (gogoproto.nullable) = false ];
+}
diff --git a/pkg/rpc/auth_tenant.go b/pkg/rpc/auth_tenant.go
@@ -115,6 +115,9 @@ func (a tenantAuthorizer) authorize(
 	case "/cockroach.server.serverpb.Status/TransactionContentionEvents":
 		return a.authTenant(tenID)
 
+	case "/cockroach.server.serverpb.Status/SpanStats":
+		return a.authSpanStats(tenID, req.(*roachpb.InternalSpanStatsRequest))
+
 	case "/cockroach.roachpb.Internal/GetSpanConfigs":
 		return a.authGetSpanConfigs(tenID, req.(*roachpb.GetSpanConfigsRequest))
 
@@ -211,6 +214,12 @@ func (a tenantAuthorizer) authGetRangeDescriptors(
 	return validateSpan(tenID, args.Span)
 }
 
+func (a tenantAuthorizer) authSpanStats(
+	tenID roachpb.TenantID, args *roachpb.InternalSpanStatsRequest,
+) error {
+	return validateSpan(tenID, args.Span)
+}
+
 // authRangeLookup authorizes the provided tenant to invoke the RangeLookup RPC
 // with the provided args.
 func (a tenantAuthorizer) authRangeLookup(

diff --git a/pkg/server/BUILD.bazel b/pkg/server/BUILD.bazel
@@ -35,6 +35,7 @@ go_library(
         "init.go",
         "init_handshake.go",
         "initial_sql.go",
+        "key_visualizer_server.go",
         "listen_and_update_addrs.go",
         "load_endpoint.go",
         "loss_of_quorum.go",
@@ -113,6 +114,8 @@ go_library(
         "//pkg/kv/kvclient/kvtenant",
         "//pkg/kv/kvclient/rangefeed",
         "//pkg/kv/kvclient/rangestats",
+        "//pkg/kv/kvclient/spanstats",
+        "//pkg/kv/kvclient/spanstats/spanstatsaccessor",
         "//pkg/kv/kvprober",
         "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/allocator/allocatorimpl",
@@ -410,6 +413,7 @@ go_test(
         "server_test.go",
         "settings_cache_test.go",
         "settings_test.go",
+        "span_stats_test.go",
         "statements_test.go",
         "stats_test.go",
         "status_ext_test.go",

diff --git a/pkg/server/admin.go b/pkg/server/admin.go
@@ -1360,7 +1360,7 @@ func (s *adminServer) statsForSpan(
 							req := serverpb.SpanStatsRequest{
 								StartKey: rSpan.Key,
 								EndKey:   rSpan.EndKey,
-								NodeID:   nodeID.String(),
+								NodeID:   nodeID,
 							}
 							spanResponse, err = client.SpanStats(ctx, &req)
 						}
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		module github.com/cockroachdb/cockroach/pkg/roachpb

		go 1.19