[WIP,DNM] server,tracing: integrate on-demand profiling with CRDB tracing

pkg/BUILD.bazel

@@ -269,6 +269,7 @@ ALL_TESTS = [
     "//pkg/security/username:username_disallowed_imports_test",
     "//pkg/security/username:username_test",
     "//pkg/security:security_test",
+    "//pkg/server/backgroundprofiler/profiler:profiler_test",
     "//pkg/server/debug/goroutineui:goroutineui_test",
     "//pkg/server/debug/pprofui:pprofui_test",
     "//pkg/server/debug:debug_test",
@@ -1413,6 +1414,10 @@ GO_TARGETS = [
     "//pkg/security/username:username_test",
     "//pkg/security:security",
     "//pkg/security:security_test",
+    "//pkg/server/backgroundprofiler/profiler:profiler",
+    "//pkg/server/backgroundprofiler/profiler:profiler_test",
+    "//pkg/server/backgroundprofiler:backgroundprofiler",
+    "//pkg/server/backgroundprofiler:executiontracer",
     "//pkg/server/debug/goroutineui:goroutineui",
     "//pkg/server/debug/goroutineui:goroutineui_test",
     "//pkg/server/debug/pprofui:pprofui",
@@ -2783,6 +2788,8 @@ GET_X_DATA_TARGETS = [
     "//pkg/security/sessionrevival:get_x_data",
     "//pkg/security/username:get_x_data",
     "//pkg/server:get_x_data",
+    "//pkg/server/backgroundprofiler:get_x_data",
+    "//pkg/server/backgroundprofiler/profiler:get_x_data",
     "//pkg/server/debug:get_x_data",
     "//pkg/server/debug/goroutineui:get_x_data",
     "//pkg/server/debug/pprofui:get_x_data",

pkg/base/constants.go

@@ -46,6 +46,10 @@ const (
 	// stores profiles when the periodic CPU profile dump is enabled.
 	CPUProfileDir = "pprof_dump"
 
+	// RuntimeProfileDir is the directory name where the
+	// backgroundprofiler.Profiler stores profiles.
+	RuntimeProfileDir = "runtime_profiler"
+
 	// InflightTraceDir is the directory name where the job trace dumper stores traces
 	// when a job opts in to dumping its execution traces.
 	InflightTraceDir = "inflight_trace_dump"

pkg/base/test_server_args.go

@@ -142,6 +142,11 @@ type TestServerArgs struct {
 	// If set, a TraceDir is initialized at the provided path.
 	TraceDir string
 
+	// If set, a RuntimeProfileDir is initialized at the provided path. Runtime
+	// profiles that are collected by backgroundprofiler.Profiler during the
+	// execution of the test will be written to this directory.
+	RuntimeProfileDir string
+
 	// DisableSpanConfigs disables the use of the span configs infrastructure
 	// (in favor of the gossiped system config span). It's equivalent to setting
 	// COCKROACH_DISABLE_SPAN_CONFIGS, and is only intended for tests written
@@ -353,6 +358,13 @@ type TestTenantArgs struct {
 	// If set, this directory should be cleaned up after the test completes.
 	HeapProfileDirName string
 
+	// RuntimeProfileDirName is used to initialize the same named field on the
+	// SQLServer.BaseConfig field. It is the directory name for runtime profiles
+	// using backgroundprofiler.Profiler. If empty, no runtime profiles will be
+	// collected during the test. If set, this directory should be cleaned up
+	// after the test completes.
+	RuntimeProfileDirName string
+
 	// StartDiagnosticsReporting checks cluster.TelemetryOptOut(), and
 	// if not disabled starts the asynchronous goroutine that checks for
 	// CockroachDB upgrades and periodically reports diagnostics to

pkg/cli/log_flags.go

@@ -227,6 +227,7 @@ func setupLogging(ctx context.Context, cmd *cobra.Command, isServerCmd, applyCon
 	serverCfg.HeapProfileDirName = filepath.Join(outputDirectory, base.HeapProfileDir)
 	serverCfg.CPUProfileDirName = filepath.Join(outputDirectory, base.CPUProfileDir)
 	serverCfg.InflightTraceDirName = filepath.Join(outputDirectory, base.InflightTraceDir)
+	serverCfg.RuntimeProfileDirName = filepath.Join(outputDirectory, base.RuntimeProfileDir)
 
 	return nil
 }

pkg/gen/protobuf.bzl

@@ -49,6 +49,7 @@ PROTOBUF_SRCS = [
     "//pkg/repstream/streampb:streampb_go_proto",
     "//pkg/roachpb:roachpb_go_proto",
     "//pkg/rpc:rpc_go_proto",
+    "//pkg/server/backgroundprofiler/profiler:profiler_go_proto",
     "//pkg/server/diagnostics/diagnosticspb:diagnosticspb_go_proto",
     "//pkg/server/serverpb:serverpb_go_proto",
     "//pkg/server/status/statuspb:statuspb_go_proto",

pkg/server/BUILD.bazel

@@ -158,6 +158,7 @@ go_library(
         "//pkg/security/password",
         "//pkg/security/securityassets",
         "//pkg/security/username",
+        "//pkg/server/backgroundprofiler/profiler",
         "//pkg/server/debug",
         "//pkg/server/diagnostics",
         "//pkg/server/diagnostics/diagnosticspb",

pkg/server/backgroundprofiler/BUILD.bazel

@@ -0,0 +1,20 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "executiontracer",
+    srcs = ["executiontracer.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/server/executiontracer",
+    visibility = ["//visibility:public"],
+    deps = ["//pkg/util/protoutil"],
+)
+
+go_library(
+    name = "backgroundprofiler",
+    srcs = ["background_profiler.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/server/backgroundprofiler",
+    visibility = ["//visibility:public"],
+    deps = ["//pkg/util/protoutil"],
+)
+
+get_x_data(name = "get_x_data")

pkg/server/backgroundprofiler/background_profiler.go

@@ -0,0 +1,59 @@
+// 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 backgroundprofiler
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+)
+
+// ProfileID is a unique identifier of the operation being profiled by the
+// Profiler.
+type ProfileID int
+
+// SubscriberID is a unique identifier of the Subscriber subscribing to the
+// background profile collection.
+type SubscriberID int
+
+// IsSet returns true if the BackgroundProfiler is currently associated with a
+// profileID.
+func (r ProfileID) IsSet() bool {
+	return r != 0
+}
+
+// Subscriber is the interface that describes an object that can subscribe to
+// the background profiler.
+type Subscriber interface {
+	// LabelValue returns the value that will be used when setting the pprof
+	// labels of the Subscriber. The key of the label will always be the ProfileID
+	// thereby allowing us to identify all samples that describe the operation
+	// being profiled.
+	LabelValue() string
+	// Identifier returns the unique identifier of the Subscriber.
+	Identifier() SubscriberID
+	// ProfileID returns the unique identifier of the operation that the
+	// Subscriber is executing on behalf of.
+	ProfileID() ProfileID
+}
+
+// Profiler is the interface that exposes methods to subscribe and unsubscribe
+// from a background profiler.
+type Profiler interface {
+	// Subscribe registers the subscriber with the background profiler. This
+	// method returns a context wrapped with pprof labels along with a closure to
+	// restore the original labels of the context.
+	Subscribe(ctx context.Context, subscriber Subscriber) (context.Context, func())
+	// Unsubscribe unregisters the subscriber from the background profiler. If the
+	// subscriber is responsible for finishing the profile the method will also
+	// return metadata describing the collected profile.
+	Unsubscribe(subscriber Subscriber) (finishedProfile bool, msg protoutil.Message)
+}

pkg/server/backgroundprofiler/profiler/BUILD.bazel

@@ -0,0 +1,64 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
+
+proto_library(
+    name = "profiler_proto",
+    srcs = ["profiler.proto"],
+    strip_import_prefix = "/pkg",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_gogo_protobuf//gogoproto:gogo_proto"],
+)
+
+go_proto_library(
+    name = "profiler_go_proto",
+    compilers = ["//pkg/cmd/protoc-gen-gogoroach:protoc-gen-gogoroach_compiler"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/server/backgroundprofiler/profiler",
+    proto = ":profiler_proto",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_gogo_protobuf//gogoproto"],
+)
+
+go_library(
+    name = "profiler",
+    srcs = ["profiler.go"],
+    embed = [":profiler_go_proto"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/server/backgroundprofiler/profiler",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/server/backgroundprofiler",
+        "//pkg/server/dumpstore",
+        "//pkg/settings",
+        "//pkg/settings/cluster",
+        "//pkg/util/log",
+        "//pkg/util/pprofutil",
+        "//pkg/util/protoutil",
+        "//pkg/util/stop",
+        "//pkg/util/syncutil",
+        "//pkg/util/timeutil",
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_google_pprof//profile",
+    ],
+)
+
+go_test(
+    name = "profiler_test",
+    srcs = ["profiler_test.go"],
+    args = ["-test.timeout=295s"],
+    deps = [
+        ":profiler",
+        "//pkg/settings/cluster",
+        "//pkg/testutils",
+        "//pkg/util/ctxgroup",
+        "//pkg/util/log",
+        "//pkg/util/stop",
+        "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingpb",
+        "@com_github_gogo_protobuf//types",
+        "@com_github_google_pprof//profile",
+        "@com_github_stretchr_testify//require",
+    ],
+)
+
+get_x_data(name = "get_x_data")