Merge pull request #154468 from yuzefovich/blathers/backport-release-25.4-154393

release-25.4: sql: fix another minor bug with CHECK EXTERNAL CONNECTION

Author: yuzefovich

pkg/sql/check_external_connection.go

@@ -90,7 +90,7 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 			time.Duration(tree.MustBeDInt(nanos)),
 		))
 	}
-	n.rows = make(chan tree.Datums, int64(len(sqlInstanceIDs))*n.params.Concurrency)
+	n.rows = make(chan tree.Datums, len(sqlInstanceIDs)*getCloudCheckConcurrency(n.params))
 	rowWriter := NewCallbackResultWriter(func(ctx context.Context, row tree.Datums) error {
 		// collapse the two pairs of bytes+time to a single string rate each.
 		res := make(tree.Datums, len(row)-1)
@@ -103,13 +103,15 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 		return nil
 	})
 
-	grp := ctxgroup.WithContext(params.ctx)
-	n.execGrp = grp
-	grp.GoCtx(func(ctx context.Context) error {
+	workerStarted := make(chan struct{})
+	n.execGrp = ctxgroup.WithContext(params.ctx)
+	n.execGrp.GoCtx(func(ctx context.Context) error {
 		// Derive a separate tracing span since the planning one will be
 		// finished when the main goroutine exits from startExec.
 		ctx, span := tracing.ChildSpan(ctx, "CheckExternalConnection-execution")
 		defer span.Finish()
+		// Unblock the main goroutine after having created the tracing span.
+		close(workerStarted)
 
 		recv := MakeDistSQLReceiver(
 			ctx,
@@ -129,6 +131,14 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 		return nil
 	})
 
+	// Block until the worker goroutine has started. This allows us to guarantee
+	// that params.ctx contains a tracing span that hasn't been finished.
+	// TODO(yuzefovich): this is a bit hacky. The issue is that
+	// planNodeToRowSource has already created a new tracing span for this
+	// checkExternalConnectionNode and has updated params.ctx accordingly; then,
+	// if the query is canceled before the worker goroutine starts, the tracing
+	// span is finished, yet it will have already been captured by the ctxgroup.
+	<-workerStarted
 	return nil
 }
 

pkg/sql/cloud_check_processor.go

@@ -203,15 +203,19 @@ func newCloudCheckProcessor(
 	return p, nil
 }
 
-// Start is part of the RowSource interface.
-func (p *proc) Start(ctx context.Context) {
-	p.StartInternal(ctx, "cloudcheck.proc")
-
-	concurrency := int(p.spec.Params.Concurrency)
+func getCloudCheckConcurrency(params CloudCheckParams) int {
+	concurrency := int(params.Concurrency)
 	if concurrency < 1 {
 		concurrency = 1
 	}
+	return concurrency
+}
+
+// Start is part of the RowSource interface.
+func (p *proc) Start(ctx context.Context) {
+	p.StartInternal(ctx, "cloudcheck.proc")
 
+	concurrency := getCloudCheckConcurrency(p.spec.Params)
 	p.results = make(chan result, concurrency)
 
 	if err := p.FlowCtx.Stopper().RunAsyncTask(p.Ctx(), "cloudcheck.proc", func(ctx context.Context) {