sql: fix interaction between stmt bundles and tracing

Previously, we wouldn't generate the bundle if the verbose tracing was
already enabled on the cluster because we wouldn't call
`instrumentationHelper.Finish` where we actually generate the bundle.
This would result in empty responses for `EXPLAIN ANALYZE (DEBUG)` as
well as the requests for stmt diagnostics being stuck in "waiting"
state.

Release note (bug fix): Previously, if the tracing
(`sql.trace.txn.enable_threshold` cluster setting) was enabled on the
cluster, the statement diagnostics collection (`EXPLAIN ANALYZE
(DEBUG)`) wouldn't work. This is now fixed.

Release justification: low-risk fix to a long-standing bug.

pkg/sql/explain_bundle_test.go

@@ -140,6 +140,16 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
  )
  r.Exec(t, `RESET enable_insert_fast_path;`)
  })
+
+ t.Run("basic when tracing already enabled", func(t *testing.T) {
+ r.Exec(t, "SET CLUSTER SETTING sql.trace.txn.enable_threshold='100ms';")
+ defer r.Exec(t, "SET CLUSTER SETTING sql.trace.txn.enable_threshold='0ms';")
+ rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM abc WHERE c=1")
+ checkBundle(
+ t, fmt.Sprint(rows), "public.abc",
+ base, plans, "stats-defaultdb.public.abc.sql", "distsql.html vec.txt vec-v.txt",
+ )
+ })
 }
 
 // checkBundle searches text strings for a bundle URL and then verifies that the

pkg/sql/instrumentation.go

@@ -95,9 +95,12 @@ type instrumentationHelper struct {
  finishCollectionDiagnostics func()
  withStatementTrace func(trace tracing.Recording, stmt string)
 
- sp *tracing.Span
- origCtx context.Context
- evalCtx *tree.EvalContext
+ sp *tracing.Span
+ // shouldFinishSpan determines whether sp needs to be finished in
+ // instrumentationHelper.Finish.
+ shouldFinishSpan bool
+ origCtx context.Context
+ evalCtx *tree.EvalContext
 
  // If savePlanForStats is true, the explainPlan will be collected and returned
  // via PlanForStats().
@@ -148,6 +151,7 @@ func (ih *instrumentationHelper) Setup(
  ih.fingerprint = fingerprint
  ih.implicitTxn = implicitTxn
  ih.codec = cfg.Codec
+ ih.origCtx = ctx
 
  switch ih.outputMode {
  case explainAnalyzeDebugOutput:
@@ -177,7 +181,11 @@ func (ih *instrumentationHelper) Setup(
  // collection is enabled so that stats are shown in the traces, but
  // no extra work is needed by the instrumentationHelper.
  ih.collectExecStats = true
- return ctx, false
+ // We still, however, want to finish the instrumentationHelper in
+ // case we're collecting a bundle. We also capture the span in order
+ // to fetch the trace from it, but the span won't be finished.
+ ih.sp = sp
+ return ctx, ih.collectBundle
  }
  } else {
  if util.CrdbTestBuild {
@@ -200,18 +208,18 @@ func (ih *instrumentationHelper) Setup(
  if ih.collectExecStats {
  // If we need to collect stats, create a non-verbose child span. Stats
  // will be added as structured metadata and processed in Finish.
- ih.origCtx = ctx
  newCtx, ih.sp = tracing.EnsureChildSpan(ctx, cfg.AmbientCtx.Tracer, "traced statement", tracing.WithForceRealSpan())
+ ih.shouldFinishSpan = true
  return newCtx, true
  }
  return ctx, false
  }
 
  ih.collectExecStats = true
  ih.traceMetadata = make(execNodeTraceMetadata)
- ih.origCtx = ctx
  ih.evalCtx = p.EvalContext()
  newCtx, ih.sp = tracing.StartVerboseTrace(ctx, cfg.AmbientCtx.Tracer, "traced statement")
+ ih.shouldFinishSpan = true
  return newCtx, true
 }
 
@@ -230,7 +238,9 @@ func (ih *instrumentationHelper) Finish(
  if ih.sp == nil {
  return retErr
  }
- ih.sp.Finish()
+ if ih.shouldFinishSpan {
+ ih.sp.Finish()
+ }
 
  // Record the statement information that we've collected.
  // Note that in case of implicit transactions, the trace contains the auto-commit too.