sql: Add new generator crdb_internal.scan_storage_internal_keys()

This new builtin is used to gather specific pebble metrics for a node
and store id (within an given keyspan). The builtin returns information
about the different types of keys (including snapshot pinned keys) as
well as bytes.

Informs: cockroachdb#94659
Release note: None

Author: craig[bot]

pkg/kv/kvserver/api.proto

@@ -89,6 +89,18 @@ message GetTableMetricsResponse{
   repeated storage.enginepb.SSTableMetricsInfo table_metrics = 1 [(gogoproto.nullable) = false, (gogoproto.customname) = "TableMetrics"];
 }
 
+// ScanStorageInternalKeysRequest retrieves metrics about keys within a range belonging
+// to a particular node and store.
+message ScanStorageInternalKeysRequest {
+  StoreRequestHeader header = 1 [(gogoproto.nullable) = false, (gogoproto.embed) = true];
+  roachpb.Span span = 2 [(gogoproto.nullable) = false];
+  int64 MegabytesPerSecond = 3;
+}
+
+message ScanStorageInternalKeysResponse {
+  repeated storage.enginepb.StorageInternalKeysMetrics advanced_metrics = 1 [(gogoproto.nullable) = false, (gogoproto.customname) = "AdvancedPebbleMetrics"];
+}
+
 message CompactEngineSpanResponse {
 }
 

pkg/kv/kvserver/storage_engine_client.go

@@ -76,6 +76,33 @@ func (c *StorageEngineClient) GetTableMetrics(
 	return resp.TableMetrics, nil
 }
 
+// ScanStorageInternalKeys is a tree.ScanStorageInternalKeys
+func (c *StorageEngineClient) ScanStorageInternalKeys(
+	ctx context.Context, nodeID, storeID int32, startKey, endKey []byte, megabytesPerSecond int64,
+) ([]enginepb.StorageInternalKeysMetrics, error) {
+	conn, err := c.nd.Dial(ctx, roachpb.NodeID(nodeID), rpc.DefaultClass)
+	if err != nil {
+		return []enginepb.StorageInternalKeysMetrics{}, errors.Wrapf(err, "could not dial node ID %d", nodeID)
+	}
+
+	client := NewPerStoreClient(conn)
+	req := &ScanStorageInternalKeysRequest{
+		StoreRequestHeader: StoreRequestHeader{
+			NodeID:  roachpb.NodeID(nodeID),
+			StoreID: roachpb.StoreID(storeID),
+		},
+		Span:               roachpb.Span{Key: roachpb.Key(startKey), EndKey: roachpb.Key(endKey)},
+		MegabytesPerSecond: megabytesPerSecond,
+	}
+
+	resp, err := client.ScanStorageInternalKeys(ctx, req)
+
+	if err != nil {
+		return []enginepb.StorageInternalKeysMetrics{}, err
+	}
+	return resp.AdvancedPebbleMetrics, nil
+}
+
 // SetCompactionConcurrency is a tree.CompactionConcurrencyFunc.
 func (c *StorageEngineClient) SetCompactionConcurrency(
 	ctx context.Context, nodeID, storeID int32, compactionConcurrency uint64,

pkg/kv/kvserver/storage_services.proto

@@ -44,5 +44,6 @@ service PerReplica {
 service PerStore {
     rpc CompactEngineSpan(cockroach.kv.kvserver.CompactEngineSpanRequest) returns (cockroach.kv.kvserver.CompactEngineSpanResponse) {}
     rpc GetTableMetrics(cockroach.kv.kvserver.GetTableMetricsRequest) returns (cockroach.kv.kvserver.GetTableMetricsResponse) {}
+    rpc ScanStorageInternalKeys(cockroach.kv.kvserver.ScanStorageInternalKeysRequest) returns (cockroach.kv.kvserver.ScanStorageInternalKeysResponse) {}
     rpc SetCompactionConcurrency(cockroach.kv.kvserver.CompactionConcurrencyRequest) returns (cockroach.kv.kvserver.CompactionConcurrencyResponse) {}
 }

pkg/kv/kvserver/stores_server.go

@@ -173,6 +173,24 @@ func (is Server) GetTableMetrics(
 	return resp, err
 }
 
+func (is Server) ScanStorageInternalKeys(
+	ctx context.Context, req *ScanStorageInternalKeysRequest,
+) (*ScanStorageInternalKeysResponse, error) {
+	resp := &ScanStorageInternalKeysResponse{}
+	err := is.execStoreCommand(ctx, req.StoreRequestHeader,
+		func(ctx context.Context, s *Store) error {
+			metrics, err := s.TODOEngine().ScanStorageInternalKeys(req.Span.Key, req.Span.EndKey, req.MegabytesPerSecond)
+
+			if err != nil {
+				return err
+			}
+
+			resp.AdvancedPebbleMetrics = metrics
+			return nil
+		})
+	return resp, err
+}
+
 // SetCompactionConcurrency implements PerStoreServer. It changes the compaction
 // concurrency of a store. While SetCompactionConcurrency is safe for concurrent
 // use, it adds uncertainty about the compaction concurrency actually set on

pkg/server/server_sql.go

@@ -961,16 +961,17 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		ClientCertExpirationCache: security.NewClientCertExpirationCache(
 			ctx, cfg.Settings, cfg.stopper, &timeutil.DefaultTimeSource{}, rootSQLMemoryMonitor,
 		),
-		RootMemoryMonitor:         rootSQLMemoryMonitor,
-		TestingKnobs:              sqlExecutorTestingKnobs,
-		CompactEngineSpanFunc:     storageEngineClient.CompactEngineSpan,
-		CompactionConcurrencyFunc: storageEngineClient.SetCompactionConcurrency,
-		GetTableMetricsFunc:       storageEngineClient.GetTableMetrics,
-		TraceCollector:            traceCollector,
-		TenantUsageServer:         cfg.tenantUsageServer,
-		KVStoresIterator:          cfg.kvStoresIterator,
-		InspectzServer:            cfg.inspectzServer,
-		RangeDescIteratorFactory:  cfg.rangeDescIteratorFactory,
+		RootMemoryMonitor:           rootSQLMemoryMonitor,
+		TestingKnobs:                sqlExecutorTestingKnobs,
+		CompactEngineSpanFunc:       storageEngineClient.CompactEngineSpan,
+		CompactionConcurrencyFunc:   storageEngineClient.SetCompactionConcurrency,
+		GetTableMetricsFunc:         storageEngineClient.GetTableMetrics,
+		ScanStorageInternalKeysFunc: storageEngineClient.ScanStorageInternalKeys,
+		TraceCollector:              traceCollector,
+		TenantUsageServer:           cfg.tenantUsageServer,
+		KVStoresIterator:            cfg.kvStoresIterator,
+		InspectzServer:              cfg.inspectzServer,
+		RangeDescIteratorFactory:    cfg.rangeDescIteratorFactory,
 		SyntheticPrivilegeCache: syntheticprivilegecache.New(
 			cfg.Settings, cfg.stopper, cfg.db,
 			serverCacheMemoryMonitor.MakeBoundAccount(),

pkg/sql/exec_util.go

@@ -1383,6 +1383,10 @@ type ExecutorConfig struct {
 	// overlap with a key range for a specified node and store.
 	GetTableMetricsFunc eval.GetTableMetricsFunc
 
+	// ScanStorageInternalKeys is used to gather information about the types of
+	// keys (including snapshot pinned keys) at each level of a node store.
+	ScanStorageInternalKeysFunc eval.ScanStorageInternalKeysFunc
+
 	// TraceCollector is used to contact all live nodes in the cluster, and
 	// collect trace spans from their inflight node registries.
 	TraceCollector *collector.TraceCollector

pkg/sql/planner.go

@@ -127,6 +127,7 @@ func (evalCtx *extendedEvalContext) copyFromExecCfg(execCfg *ExecutorConfig) {
 	evalCtx.CompactEngineSpan = execCfg.CompactEngineSpanFunc
 	evalCtx.SetCompactionConcurrency = execCfg.CompactionConcurrencyFunc
 	evalCtx.GetTableMetrics = execCfg.GetTableMetricsFunc
+	evalCtx.ScanStorageInternalKeys = execCfg.ScanStorageInternalKeysFunc
 	evalCtx.TestingKnobs = execCfg.EvalContextTestingKnobs
 	evalCtx.ClusterID = execCfg.NodeInfo.LogicalClusterID()
 	evalCtx.ClusterName = execCfg.RPCContext.ClusterName()

pkg/sql/sem/builtins/fixed_oids.go

@@ -2440,6 +2440,8 @@ var builtinOidsArray = []string{
 	2467: `crdb_internal.request_statement_bundle(stmtFingerprint: string, planGist: string, samplingProbability: float, minExecutionLatency: interval, expiresAfter: interval) -> bool`,
 	2468: `crdb_internal.request_statement_bundle(stmtFingerprint: string, planGist: string, antiPlanGist: bool, samplingProbability: float, minExecutionLatency: interval, expiresAfter: interval) -> bool`,
 	2469: `crdb_internal.is_system_table_key(raw_key: bytes) -> bool`,
+	2470: `crdb_internal.scan_storage_internal_keys(node_id: int, store_id: int, start_key: bytes, end_key: bytes) -> tuple{int AS level, int AS node_id, int AS store_id, int AS snapshot_pinned_keys, int AS snapshot_pinned_keys_bytes, int AS point_key_delete_count, int AS point_key_set_count, int AS range_delete_count, int AS range_key_set_count, int AS range_key_delete_count}`,
+	2471: `crdb_internal.scan_storage_internal_keys(node_id: int, store_id: int, start_key: bytes, end_key: bytes, mb_per_second: int4) -> tuple{int AS level, int AS node_id, int AS store_id, int AS snapshot_pinned_keys, int AS snapshot_pinned_keys_bytes, int AS point_key_delete_count, int AS point_key_set_count, int AS range_delete_count, int AS range_key_set_count, int AS range_key_delete_count}`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/builtins/generator_builtins.go

@@ -684,6 +684,36 @@ The last argument is a JSONB object containing the following optional fields:
 			volatility.Stable,
 		),
 	),
+	"crdb_internal.scan_storage_internal_keys": makeBuiltin(
+		tree.FunctionProperties{
+			Category: builtinconstants.CategorySystemInfo,
+		},
+		makeGeneratorOverload(
+			tree.ParamTypes{
+				{Name: "node_id", Typ: types.Int},
+				{Name: "store_id", Typ: types.Int},
+				{Name: "start_key", Typ: types.Bytes},
+				{Name: "end_key", Typ: types.Bytes},
+			},
+			storageInternalKeysGeneratorType,
+			makeStorageInternalKeysGenerator,
+			"Scans a store's storage engine, computing statistics describing the internal keys within the span [start_key, end_key). This function is rate limited to 10 megabytes per second.",
+			volatility.Volatile,
+		),
+		makeGeneratorOverload(
+			tree.ParamTypes{
+				{Name: "node_id", Typ: types.Int},
+				{Name: "store_id", Typ: types.Int},
+				{Name: "start_key", Typ: types.Bytes},
+				{Name: "end_key", Typ: types.Bytes},
+				{Name: "mb_per_second", Typ: types.Int4},
+			},
+			storageInternalKeysGeneratorType,
+			makeStorageInternalKeysGenerator,
+			"Scans a store's storage engine, computing statistics describing the internal keys within the span [start_key, end_key).",
+			volatility.Volatile,
+		),
+	),
 }
 
 var decodePlanGistGeneratorType = types.String
@@ -3272,6 +3302,116 @@ func makeTableMetricsGenerator(
 	return newTableMetricsIterator(evalCtx, nodeID, storeID, start, end), nil
 }
 
+type storageInternalKeysIterator struct {
+	metrics []enginepb.StorageInternalKeysMetrics
+	evalCtx *eval.Context
+
+	iterIdx            int
+	nodeID             int32
+	storeID            int32
+	megabytesPerSecond int64
+	start              []byte
+	end                []byte
+}
+
+var storageInternalKeysGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.Int, types.Int, types.Int, types.Int, types.Int, types.Int, types.Int, types.Int,
+		types.Int, types.Int},
+	[]string{
+		"level",
+		"node_id",
+		"store_id",
+		"snapshot_pinned_keys",
+		"snapshot_pinned_keys_bytes",
+		"point_key_delete_count",
+		"point_key_set_count",
+		"range_delete_count",
+		"range_key_set_count",
+		"range_key_delete_count",
+	},
+)
+
+var _ eval.ValueGenerator = (*storageInternalKeysIterator)(nil)
+
+func newStorageInternalKeysGenerator(
+	evalCtx *eval.Context, nodeID, storeID int32, start, end []byte, megaBytesPerSecond int64,
+) *storageInternalKeysIterator {
+	return &storageInternalKeysIterator{evalCtx: evalCtx, nodeID: nodeID, storeID: storeID, start: start, end: end, megabytesPerSecond: megaBytesPerSecond}
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (s *storageInternalKeysIterator) Start(ctx context.Context, _ *kv.Txn) error {
+	var err error
+	s.metrics, err = s.evalCtx.ScanStorageInternalKeys(ctx, s.nodeID, s.storeID, s.start, s.end, s.megabytesPerSecond)
+	if err != nil {
+		err = errors.Wrapf(err, "getting table metrics for node %d store %d", s.nodeID, s.storeID)
+	}
+
+	return err
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (s *storageInternalKeysIterator) Next(_ context.Context) (bool, error) {
+	s.iterIdx++
+	return s.iterIdx <= len(s.metrics), nil
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (s *storageInternalKeysIterator) Values() (tree.Datums, error) {
+	metricsInfo := s.metrics[s.iterIdx-1]
+	levelDatum := tree.DNull
+
+	if metricsInfo.Level != -1 {
+		levelDatum = tree.NewDInt(tree.DInt(metricsInfo.Level))
+	}
+
+	return tree.Datums{
+		levelDatum,
+		tree.NewDInt(tree.DInt(s.nodeID)),
+		tree.NewDInt(tree.DInt(s.storeID)),
+		tree.NewDInt(tree.DInt(metricsInfo.SnapshotPinnedKeys)),
+		tree.NewDInt(tree.DInt(metricsInfo.SnapshotPinnedKeysBytes)),
+		tree.NewDInt(tree.DInt(metricsInfo.PointKeyDeleteCount)),
+		tree.NewDInt(tree.DInt(metricsInfo.PointKeySetCount)),
+		tree.NewDInt(tree.DInt(metricsInfo.RangeDeleteCount)),
+		tree.NewDInt(tree.DInt(metricsInfo.RangeKeySetCount)),
+		tree.NewDInt(tree.DInt(metricsInfo.RangeKeyDeleteCount)),
+	}, nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (tmi *storageInternalKeysIterator) Close(_ context.Context) {}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (tmi *storageInternalKeysIterator) ResolvedType() *types.T {
+	return storageInternalKeysGeneratorType
+}
+
+func makeStorageInternalKeysGenerator(
+	ctx context.Context, evalCtx *eval.Context, args tree.Datums,
+) (eval.ValueGenerator, error) {
+	isAdmin, err := evalCtx.SessionAccessor.HasAdminRole(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if !isAdmin {
+		return nil, errInsufficientPriv
+	}
+	nodeID := int32(tree.MustBeDInt(args[0]))
+	storeID := int32(tree.MustBeDInt(args[1]))
+	start := []byte(tree.MustBeDBytes(args[2]))
+	end := []byte(tree.MustBeDBytes(args[3]))
+
+	var megabytesPerSecond int64
+	if len(args) > 4 {
+		megabytesPerSecond = int64(tree.MustBeDInt(args[4]))
+	} else {
+		megabytesPerSecond = int64(10)
+	}
+
+	return newStorageInternalKeysGenerator(evalCtx, nodeID, storeID, start, end, megabytesPerSecond), nil
+}
+
 var tableSpanStatsGeneratorType = types.MakeLabeledTuple(
 	[]*types.T{types.Int, types.Int, types.Int, types.Int, types.Int, types.Int, types.Float},
 	[]string{"database_id", "table_id", "range_count", "approximate_disk_bytes", "live_bytes", "total_bytes", "live_percentage"},

pkg/sql/sem/eval/context.go

@@ -227,6 +227,9 @@ type Context struct {
 	// GetTableMetrics is used in crdb_internal.sstable_metrics.
 	GetTableMetrics GetTableMetricsFunc
 
+	// ScanStorageInternalKeys is used in crdb_internal.scan_storage_internal_keys.
+	ScanStorageInternalKeys ScanStorageInternalKeysFunc
+
 	// SetCompactionConcurrency is used to change the compaction concurrency of
 	// a store.
 	SetCompactionConcurrency SetCompactionConcurrencyFunc