Merge #158931 #159113

craig[bot] · dhartunian · yuzefovich · craig[bot] · commit 1645953cf0c3 · 2025-12-10T17:17:28.000Z
158931: obs, kv, sql: add workloadID to profile tags r=tbg,jasonlmfong a=dhartunian 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 159113: kvcoord: optimize txn write buffer for read-only txns r=yuzefovich a=yuzefovich In the txn write buffer we need to merge the ScanResponse coming from the server with any overlapping writes that we've buffered. This requires us to allocate `respMerger.batchResponses` which is the "staging" area of the merged response. However, if we don't have any overlapping buffered writes, then we simply copy the server response into that staging area. This commit adds an optimization to avoid this extra copy in a special but common case when the buffer is empty (i.e. we're in a read-only txn). We could extend the optimization further but that is left as a TODO. Epic: None Release note: None Co-authored-by: David Hartunian <davidh@cockroachlabs.com> Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>
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/kvcoord/txn_interceptor_write_buffer.go b/pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go
@@ -1162,6 +1162,15 @@ func (twb *txnWriteBuffer) mergeWithScanResp(
 			"with COL_BATCH_RESPONSE scan format")
 	}
 
+	if twb.buffer.Len() == 0 {
+		// If we haven't buffered any writes, then we can just return the server
+		// response unchanged.
+		// TODO(yuzefovich): we could take the optimization further by examining
+		// whether any buffered writes overlap with the Scan request and
+		// skipping the merge step if not.
+		return resp, nil
+	}
+
 	respIter := newScanRespIter(req, resp)
 	// First, calculate the size of the merged response. This then allows us to
 	// exactly pre-allocate the response slice when constructing the respMerger.
@@ -1186,6 +1195,15 @@ func (twb *txnWriteBuffer) mergeWithReverseScanResp(
 			"ReverseScanRequest with COL_BATCH_RESPONSE scan format")
 	}
 
+	if twb.buffer.Len() == 0 {
+		// If we haven't buffered any writes, then we can just return the server
+		// response unchanged.
+		// TODO(yuzefovich): we could take the optimization further by examining
+		// whether any buffered writes overlap with the ReverseScan request and
+		// skipping the merge step if not.
+		return resp, nil
+	}
+
 	respIter := newReverseScanRespIter(req, resp)
 	// First, calculate the size of the merged response. This then allows us to
 	// exactly pre-allocate the response slice when constructing the respMerger.
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(),
