release-22.1: backport distsql profiler labels for remote flows

Backport of cockroachdb#92775 Fixes: cockroachdb#82464 Release note: SQL queries running on remote notes now show up in cpu profiles with "distsql.stmt" label.
cucaroach · Feb 13, 2023 · cdd538a · cdd538a
1 parent 96481c6
commit cdd538a
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/pkg/sql/distsql/BUILD.bazel b/pkg/sql/distsql/BUILD.bazel
@@ -28,6 +28,7 @@ go_library(
         "//pkg/util/envutil",
         "//pkg/util/log",
         "//pkg/util/mon",
+        "//pkg/util/pprofutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_cockroachdb_errors//:errors",

diff --git a/pkg/sql/distsql/server.go b/pkg/sql/distsql/server.go
@@ -38,6 +38,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/cockroachdb/cockroach/pkg/util/pprofutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
@@ -408,6 +409,16 @@ func (ds *ServerImpl) setupFlow(
 
 	if !f.IsLocal() {
 		flowCtx.AmbientContext.AddLogTag("f", f.GetFlowCtx().ID.Short())
+		if req.StatementSQL != "" {
+			flowCtx.AmbientContext.AddLogTag("distsql.stmt", req.StatementSQL)
+		}
+		flowCtx.AmbientContext.AddLogTag("distsql.gateway", req.Flow.Gateway)
+		if req.EvalContext.SessionData.ApplicationName != "" {
+			flowCtx.AmbientContext.AddLogTag("distsql.appname", req.EvalContext.SessionData.ApplicationName)
+		}
+		if leafTxn != nil {
+			flowCtx.AmbientContext.AddLogTag("distsql.txn", leafTxn.ID())
+		}
 		ctx = flowCtx.AmbientContext.AnnotateCtx(ctx)
 		telemetry.Inc(sqltelemetry.DistSQLExecCounter)
 	}
@@ -644,6 +655,9 @@ func (ds *ServerImpl) SetupFlow(
 		nil /* batchSyncFlowConsumer */, LocalState{},
 	)
 	if err == nil {
+		var undo func()
+		ctx, undo = pprofutil.SetProfilerLabelsFromCtxTags(ctx)
+		defer undo()
 		err = ds.flowScheduler.ScheduleFlow(ctx, f)
 	}
 	if err != nil {
@@ -690,7 +704,10 @@ func (ds *ServerImpl) flowStreamInt(
 	}
 	defer cleanup()
 	log.VEventf(ctx, 1, "connected inbound stream %s/%d", flowID.Short(), streamID)
-	return streamStrategy.Run(f.AmbientContext.AnnotateCtx(ctx), stream, msg, f)
+	ctx = f.AmbientContext.AnnotateCtx(ctx)
+	ctx, undo := pprofutil.SetProfilerLabelsFromCtxTags(ctx)
+	defer undo()
+	return streamStrategy.Run(ctx, stream, msg, f)
 }
 
 // FlowStream is part of the execinfrapb.DistSQLServer interface.

diff --git a/pkg/util/pprofutil/labels.go b/pkg/util/pprofutil/labels.go
@@ -37,7 +37,7 @@ func SetProfilerLabels(ctx context.Context, labels ...string) (_ context.Context
 func SetProfilerLabelsFromCtxTags(ctx context.Context) (_ context.Context, undo func()) {
 	tags := logtags.FromContext(ctx)
 	if tags == nil || len(tags.Get()) == 0 {
-		return
+		return ctx, func() {}
 	}
 	var labels []string
 	for _, tag := range tags.Get() {