sql: fix another minor bug with CHECK EXTERNAL CONNECTION

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.

Release note: None

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