obs, kv, sql: add workloadID to profile tags

dhartunian · dhartunian · commit c399922c40b0 · 2025-12-01T18:53:20.000Z
We introduce `pkg/obs/workloadid` which defines a commond workloadID definition, subject to extensive future modification, that we can use to tag CPU profiles and execution trace regions. Currently, the workloadID is either a static value, or a statement fingerprint ID. Where appropriate, we've tagged CPU profiles with a workload identifier. When tracing is enabled, on a request, we will attempt to include the workloadID into the region tag in the execution trace. This path is a bit perf sensitive so we only want to incur the cost of allocating the extra string when we're already paying the cost of tracing. It is expected that the combination of statementb diagnostic bundle capture and execution tracing can then be cross-referenced. Epic: CRDB-55080 Resolves: CRDB-55928 Resolves: CRDB-55924 Release note: None
diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel
@@ -1681,6 +1681,7 @@ GO_TARGETS = [
     "//pkg/obs/eventagg:eventagg_test",
     "//pkg/obs/logstream:logstream",
     "//pkg/obs/logstream:logstream_test",
+    "//pkg/obs/workloadid:workloadid",
     "//pkg/obs:obs",
     "//pkg/raft/confchange:confchange",
     "//pkg/raft/confchange:confchange_test",
diff --git a/pkg/crosscluster/logical/BUILD.bazel b/pkg/crosscluster/logical/BUILD.bazel
@@ -46,6 +46,7 @@ go_library(
         "//pkg/kv",
         "//pkg/kv/bulk",
         "//pkg/kv/kvpb",
+        "//pkg/obs/workloadid",
         "//pkg/repstream/streampb",
         "//pkg/roachpb",
         "//pkg/security/username",
diff --git a/pkg/crosscluster/logical/logical_replication_writer_processor.go b/pkg/crosscluster/logical/logical_replication_writer_processor.go
@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	kvbulk "github.com/cockroachdb/cockroach/pkg/kv/bulk"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -359,7 +360,8 @@ func (lrw *logicalReplicationWriterProcessor) Start(ctx context.Context) {
 				log.Dev.Infof(lrw.Ctx(), "consumer completed. Error: %s", err)
 				lrw.sendError(errors.Wrap(err, "consume events"))
 			}
-		}, "proc", fmt.Sprintf("%d", lrw.ProcessorID))
+		}, workloadid.ProfileTag, workloadid.WORKLOAD_NAME_LDR,
+			"proc", fmt.Sprintf("%d", lrw.ProcessorID))
 		return nil
 	})
 }
diff --git a/pkg/crosscluster/logical/offline_initial_scan_processor.go b/pkg/crosscluster/logical/offline_initial_scan_processor.go
@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/crosscluster/streamclient"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/kv/bulk"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
@@ -235,7 +236,8 @@ func (o *offlineInitialScanProcessor) Start(ctx context.Context) {
 			if err := o.subscription.Err(); err != nil {
 				o.sendError(errors.Wrap(err, "subscription"))
 			}
-		}, "proc", fmt.Sprintf("%d", o.ProcessorID))
+		}, workloadid.ProfileTag, workloadid.WORKLOAD_NAME_LDR,
+			"proc", fmt.Sprintf("%d", o.ProcessorID))
 		return nil
 	})
 }
diff --git a/pkg/kv/kvclient/rangefeed/rangefeed.go b/pkg/kv/kvclient/rangefeed/rangefeed.go
@@ -241,7 +241,9 @@ func (f *RangeFeed) start(
 		// pprof.Do function does exactly what we do here, but it also results in
 		// pprof.Do function showing up in the stack traces -- so, just set and reset
 		// labels manually.
-		ctx, reset := pprofutil.SetProfilerLabels(ctx, append(f.extraPProfLabels, "rangefeed", f.name)...)
+		ctx, reset := pprofutil.SetProfilerLabels(
+			ctx, append(f.extraPProfLabels, "rangefeed", f.name)...,
+		)
 		defer reset()
 
 		if f.invoker != nil {
diff --git a/pkg/kv/kvserver/BUILD.bazel b/pkg/kv/kvserver/BUILD.bazel
@@ -191,6 +191,7 @@ go_library(
         "//pkg/multitenant",
         "//pkg/multitenant/tenantcapabilities/tenantcapabilitiesauthorizer",
         "//pkg/multitenant/tenantcostmodel",
+        "//pkg/obs/workloadid",
         "//pkg/raft",
         "//pkg/raft/raftlogger",
         "//pkg/raft/raftpb",
diff --git a/pkg/kv/kvserver/raft_transport.go b/pkg/kv/kvserver/raft_transport.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol/kvflowcontrolpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol/node_rac2"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
@@ -931,7 +932,7 @@ func (t *RaftTransport) startProcessNewQueue(
 	}
 	go func(ctx context.Context) {
 		defer hdl.Activate(ctx).Release(ctx)
-		pprofutil.Do(ctx, worker, "remote_node_id", toNodeID.String())
+		pprofutil.Do(ctx, worker, workloadid.ProfileTag, workloadid.WORKLOAD_NAME_RAFT, "remote_node_id", toNodeID.String())
 	}(ctx)
 	return true
 }
diff --git a/pkg/kv/kvserver/replica_send.go b/pkg/kv/kvserver/replica_send.go
@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/replicastats"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/txnwait"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -139,7 +140,21 @@ func (r *Replica) SendWithWriteBytes(
 		}
 		defer reset()
 	}
+
 	if trace.IsEnabled() {
+		foundLabel := ""
+		for i, l := range ba.ProfileLabels {
+			if i%2 == 0 && l == workloadid.ProfileTag && i < len(ba.ProfileLabels)-1 {
+				// This label is set in conn_executor_exec if tracing is active.
+				foundLabel = ba.ProfileLabels[i+1]
+				break
+			}
+		}
+		// This construction avoids calling `defer` in a loop which is
+		// not permitted by our linter.
+		if foundLabel != "" {
+			defer trace.StartRegion(ctx, foundLabel).End()
+		}
 		defer trace.StartRegion(ctx, r.rangeStr.String() /* cheap */).End()
 	}
 	// Add the range log tag.
diff --git a/pkg/kv/kvserver/storeliveness/BUILD.bazel b/pkg/kv/kvserver/storeliveness/BUILD.bazel
@@ -21,6 +21,7 @@ go_library(
         "//pkg/base",
         "//pkg/keys",
         "//pkg/kv/kvserver/storeliveness/storelivenesspb",
+        "//pkg/obs/workloadid",
         "//pkg/roachpb",
         "//pkg/rpc/nodedialer",
         "//pkg/rpc/rpcbase",
diff --git a/pkg/kv/kvserver/storeliveness/transport.go b/pkg/kv/kvserver/storeliveness/transport.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	slpb "github.com/cockroachdb/cockroach/pkg/kv/kvserver/storeliveness/storelivenesspb"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
 	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
@@ -566,7 +567,10 @@ func (t *Transport) startProcessNewQueue(
 	err := t.stopper.RunAsyncTask(
 		ctx, "storeliveness.Transport: sending messages",
 		func(ctx context.Context) {
-			pprofutil.Do(ctx, worker, "remote_node_id", toNodeID.String())
+			pprofutil.Do(ctx, worker,
+				workloadid.ProfileTag, workloadid.WORKLOAD_NAME_STORELIVENESS,
+				"remote_node_id", toNodeID.String(),
+			)
 		},
 	)
 	if err != nil {
diff --git a/pkg/obs/workloadid/BUILD.bazel b/pkg/obs/workloadid/BUILD.bazel
@@ -0,0 +1,8 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "workloadid",
+    srcs = ["workloadid.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/obs/workloadid",
+    visibility = ["//visibility:public"],
+)
diff --git a/pkg/obs/workloadid/workloadid.go b/pkg/obs/workloadid/workloadid.go
@@ -0,0 +1,51 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package workloadid
+
+// This package defines CRDB-wide workload identifiers that we can
+// define and use to attribute observability data.
+//
+// We will prefer using uint64 workload identifiers because those are
+// cheaper to move around the system and can help us avoid allocation
+// in hot paths.
+//
+// String-based identifiers are useful when we want to display the
+// identifier to a human, or to label go profiles and execution traces,
+// which only accept string values for tags.
+//
+// Currently, this definition is quite simple consisting only of a
+// preset list of static identifiers below that can be used numerically
+// or using static strings. Statement Fingerprint IDs are the only
+// example here of a dynamic ID that can have high cardinality.
+//
+// It is expected that future work will expand the structure and
+// expressivity of the workloadID. No ID stability is guaranteed.
+// While we currently do persist statement fingerprint IDs in the SQL
+// Activity tables, we advise users of these values to AVOID PERSISTING
+// WORKLOAD IDs until they are stable.
+
+const ProfileTag = "workload.id"
+
+type WorkloadID uint64
+
+// The IDs and Names below are currently not persisted anywhere and
+// are subject to change.
+
+const (
+	WORKLOAD_ID_UNKNOWN WorkloadID = iota
+	WORKLOAD_ID_LDR
+	WORKLOAD_ID_RAFT
+	WORKLOAD_ID_STORELIVENESS
+	WORKLOAD_ID_RPC_HEARTBEAT
+)
+
+const (
+	WORKLOAD_NAME_UNKNOWN       = "UNKNOWN"
+	WORKLOAD_NAME_LDR           = "LDR"
+	WORKLOAD_NAME_RAFT          = "RAFT"
+	WORKLOAD_NAME_STORELIVENESS = "STORELIVENESS"
+	WORKLOAD_NAME_RPC_HEARTBEAT = "RPC_HEARTBEAT"
+)
diff --git a/pkg/rpc/BUILD.bazel b/pkg/rpc/BUILD.bazel
@@ -36,6 +36,7 @@ go_library(
         "//pkg/keys",
         "//pkg/kv/kvpb",
         "//pkg/multitenant/tenantcapabilities",
+        "//pkg/obs/workloadid",
         "//pkg/roachpb",
         "//pkg/rpc/rpcbase",
         "//pkg/security",
diff --git a/pkg/rpc/peer.go b/pkg/rpc/peer.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/VividCortex/ewma"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/circuit"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
@@ -271,7 +272,9 @@ func newPeer[Conn rpcConn](rpcCtx *Context, k peerKey, peerOpts *peerOptions[Con
 		AsyncProbe: func(report func(error), done func()) {
 			pprofutil.Do(ctx, func(ctx context.Context) {
 				p.launch(ctx, report, done)
-			}, "tags", logtags.FromContext(ctx).String())
+			},
+				workloadid.ProfileTag, workloadid.WORKLOAD_NAME_RPC_HEARTBEAT,
+				"tags", logtags.FromContext(ctx).String())
 		},
 	})
 	p.b = b
diff --git a/pkg/sql/BUILD.bazel b/pkg/sql/BUILD.bazel
@@ -350,6 +350,7 @@ go_library(
         "//pkg/multitenant/multitenantcpu",
         "//pkg/multitenant/tenantcapabilities",
         "//pkg/multitenant/tenantcapabilitiespb",
+        "//pkg/obs/workloadid",
         "//pkg/repstream",
         "//pkg/repstream/streampb",
         "//pkg/roachpb",
diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
 	"github.com/cockroachdb/cockroach/pkg/multitenant/multitenantcpu"
+	"github.com/cockroachdb/cockroach/pkg/obs/workloadid"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -49,6 +50,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessionphase"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats/sqlstatsutil"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
@@ -4438,20 +4440,27 @@ func (ex *connExecutor) execWithProfiling(
 	ctx context.Context, ast tree.Statement, prepared *prep.Statement, op func(context.Context) error,
 ) error {
 	var err error
-	if ex.server.cfg.Settings.CPUProfileType() == cluster.CPUProfileWithLabels {
+	if ex.server.cfg.Settings.CPUProfileType() == cluster.CPUProfileWithLabels || log.HasSpan(ctx) {
 		remoteAddr := "internal"
 		if rAddr := ex.sessionData().RemoteAddr; rAddr != nil {
 			remoteAddr = rAddr.String()
 		}
+		// Compute stmtNoConstants with proper FmtFlags for consistency with
+		// makeStatement and ih.Setup.
+		fmtFlags := tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&ex.server.cfg.Settings.SV))
 		var stmtNoConstants string
 		if prepared != nil {
 			stmtNoConstants = prepared.StatementNoConstants
 		} else {
-			stmtNoConstants = tree.FormatStatementHideConstants(ast)
+			stmtNoConstants = tree.FormatStatementHideConstants(ast, fmtFlags)
 		}
+		// Compute fingerprint ID here since ih.Setup hasn't been called yet.
+		fingerprintID := appstatspb.ConstructStatementFingerprintID(
+			stmtNoConstants, ex.implicitTxn(), ex.sessionData().Database)
 		pprofutil.Do(ctx, func(ctx context.Context) {
 			err = op(ctx)
 		},
+			workloadid.ProfileTag, sqlstatsutil.EncodeStmtFingerprintIDToString(fingerprintID),
 			"appname", ex.sessionData().ApplicationName,
 			"addr", remoteAddr,
 			"stmt.tag", ast.StatementTag(),
