distsql: harden infrastructure in some edge cases

As part of looking into a nil pointer crash when rows produced by the
EXPORT were used as an input to the mutation (which is fixed in the
following commit), I noticed some problems with the distsql
infrastructure. In particular, in
f397730 we didn't properly implement
`Close` method for some processors (two EXPORT ones and the column
backfiller) which would mean possibly incomplete cleanup when the
processor never runs.

Additionally, the root cause of the panic was `nil` txn that previously
silently could've been set at the very end of the flow setup if
only at that point we decided that we need to create a leaf txn, yet
LeafTxnInputState was nil. That condition is expected in some cases
(when the flow is not running under a txn, like some bulk flows), yet
for EXPORT which runs under a txn, this was unexpected. This commit adds
an assertion for this.

Furthermore, if we hit an error at this point, we need to perform
cleanup differently from earlier error paths since we've already fully
set up the flow. This commit also adds that.

Release note: None

Author: yuzefovich

pkg/sql/backfill/backfill.go

@@ -280,9 +280,12 @@ func (cb *ColumnBackfiller) InitForDistributedUse(
 
 // Close frees the resources used by the ColumnBackfiller.
 func (cb *ColumnBackfiller) Close(ctx context.Context) {
-	cb.fetcher.Close(ctx)
 	if cb.mon != nil {
+		// fetcher is only initialized when mon has been set. If mon is nil,
+		// then Close has already been called.
+		cb.fetcher.Close(ctx)
 		cb.mon.Stop(ctx)
+		cb.mon = nil
 	}
 }
 

pkg/sql/distsql/server.go

@@ -200,26 +200,29 @@ func (ds *ServerImpl) setupFlow(
 	var sp *tracing.Span                       // will be Finish()ed by Flow.Cleanup()
 	var monitor, diskMonitor *mon.BytesMonitor // will be closed in Flow.Cleanup()
 	var onFlowCleanupEnd func(context.Context) // will be called at the very end of Flow.Cleanup()
+	var cleanupPerformed bool
 	// Make sure that we clean up all resources (which in the happy case are
 	// cleaned up in Flow.Cleanup()) if an error is encountered.
 	defer func() {
 		if retErr != nil {
-			if monitor != nil {
-				monitor.Stop(ctx)
-			}
-			if diskMonitor != nil {
-				diskMonitor.Stop(ctx)
-			}
-			if onFlowCleanupEnd != nil {
-				onFlowCleanupEnd(ctx)
-			} else {
-				reserved.Close(ctx)
-				onFlowCleanup.Do()
-			}
-			// We finish the span after performing other cleanup in case that
-			// cleanup accesses the context with the span.
-			if sp != nil {
-				sp.Finish()
+			if !cleanupPerformed {
+				if monitor != nil {
+					monitor.Stop(ctx)
+				}
+				if diskMonitor != nil {
+					diskMonitor.Stop(ctx)
+				}
+				if onFlowCleanupEnd != nil {
+					onFlowCleanupEnd(ctx)
+				} else {
+					reserved.Close(ctx)
+					onFlowCleanup.Do()
+				}
+				// We finish the span after performing other cleanup in case that
+				// cleanup accesses the context with the span.
+				if sp != nil {
+					sp.Finish()
+				}
 			}
 			retCtx = tracing.ContextWithSpan(ctx, nil)
 		}
@@ -267,6 +270,9 @@ func (ds *ServerImpl) setupFlow(
 	makeLeaf := func(ctx context.Context) (*kv.Txn, error) {
 		tis := req.LeafTxnInputState
 		if tis == nil {
+			if localState.Txn != nil {
+				return nil, errors.AssertionFailedf("nil LeafTxnInputState when trying to create the LeafTxn")
+			}
 			// This must be a flow running for some bulk-io operation that doesn't use
 			// a txn.
 			return nil, nil
@@ -447,6 +453,12 @@ func (ds *ServerImpl) setupFlow(
 		if leafTxn == nil {
 			leafTxn, err = makeLeaf(ctx)
 			if err != nil {
+				// Given that we've already fully set up the flow, we must do
+				// the full cleanup. This supersedes the cleanup done in the
+				// defer at the beginning of the method, so we mark
+				// cleanupPerformed accordingly.
+				f.Cleanup(ctx)
+				cleanupPerformed = true
 				return nil, nil, nil, err
 			}
 		}

pkg/sql/export/exportcsv.go

@@ -306,4 +306,6 @@ func (sp *csvWriter) Resume(output execinfra.RowReceiver) {
 }
 
 // Close is part of the execinfra.Processor interface.
-func (*csvWriter) Close(context.Context) {}
+func (sp *csvWriter) Close(context.Context) {
+	sp.input.ConsumerClosed()
+}

pkg/sql/export/exportparquet.go

@@ -247,7 +247,9 @@ func (sp *parquetWriterProcessor) Resume(output execinfra.RowReceiver) {
 }
 
 // Close is part of the execinfra.Processor interface.
-func (*parquetWriterProcessor) Close(context.Context) {}
+func (sp *parquetWriterProcessor) Close(context.Context) {
+	sp.input.ConsumerClosed()
+}
 
 // Resume is part of the execinfra.Processor interface.
 func (sp *parquetWriterProcessor) testingKnobsOrNil() *TestingKnobs {

pkg/sql/rowexec/columnbackfiller.go

@@ -107,7 +107,9 @@ func (*columnBackfiller) Resume(output execinfra.RowReceiver) {
 }
 
 // Close is part of the execinfra.Processor interface.
-func (*columnBackfiller) Close(context.Context) {}
+func (cb *columnBackfiller) Close(ctx context.Context) {
+	cb.ColumnBackfiller.Close(ctx)
+}
 
 func (cb *columnBackfiller) doRun(ctx context.Context) *execinfrapb.ProducerMetadata {
 	finishedSpans, err := cb.mainLoop(ctx)