Merge #154393

154393: sql: fix another minor bug with CHECK EXTERNAL CONNECTION r=yuzefovich a=yuzefovich

Previously, we could create a channel with no buffer if concurrency parameter wasn't specified. As a result, in an edge case (that was exposed by recently added test) the reader goroutine might have exited _before_ reading from the unbuffered channel which would result in a deadlock (because the writer goroutine would block on sending into the channel, and the reader - i.e. the main goroutine - would block waiting on the wait group). This is now fixed by unifying the logic for determining the concurrency.

We could've made sending on `n.rows` channel also check the context cancellation if we're being conservative, but that seems like an overkill.

Additionally, it fixes an extreme edge case where we could hit "use of Span after Finish" issue in case when the context canceled (because of small statement timeout) before the execution worker goroutine was started. This is addressed by blocking the main goroutine until the worker goroutine is spun off.

Epic: None
Release note: None

Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: craig[bot]

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) {