Merge #153951 #154123

153951: sql/export: refactor processors to implement RowSource interface r=yuzefovich a=yuzefovich

**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.

**sql/export: refactor processors to implement RowSource interface**

We recently found an edge case where the output of EXPORT was used as an
input to a mutation. Currently, in such a setup we have two goroutines
since the export processors don't implement `execinfra.RowSource`
interface, so we cannot fuse them with the mutation processor. Presence
of concurrency forces usage of the LeafTxns, yet mutations require
access to the RootTxn, thus we have conflicting requirements. Before the
previous commit we'd get a nil pointer crash, with the previous
commit in place we now get an assertion failure.

This commit solves this issue by refactoring both processors to
implement the `RowSource` interface. In the problematic query it allows
all processors to be fused together to run in a single goroutine, so we
can use the RootTxn as required by the mutation.

Release note (bug fix): Previously, EXPORT CSV and EXPORT PARQUET stmts
could result in a node crash when their result rows were used as the
input to a mutation like an INSERT within the same SQL. This bug has
been present since before 22.1 version and has now been fixed.

Fixes: #153292.

154123: drpc, server: implement filtering rules for server initialization state r=cthumuluru-crdb,shubhamdhama a=Nukitt

Previously, in DRPC, the `batch` RPC was able to access `node.Descriptor` which was yet to be initialized. This led to a data race, described in #153948.

This wasn't the case for gRPC, because of an interceptor that could filter RPCs during different server states. Particularly, it allows only bootstrap, heartbeat, health and gossip methods during initialization to prevent calls to potentially uninitialized services.

This patch adds this interceptor support into drpc as well. In-depth discussion on this: #153948 (comment)

Epic: None
Fixes: #153948
Release note: None

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

Author: 3 people

pkg/server/drpc_server.go

@@ -31,7 +31,17 @@ type drpcServer struct {
 // newDRPCServer creates and configures a new drpcServer instance. It enables
 // DRPC if the experimental setting is on, otherwise returns a dummy server.
 func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error) {
-	d, err := rpc.NewDRPCServer(ctx, rpcCtx)
+	drpcServer := &drpcServer{}
+	drpcServer.setMode(modeInitializing)
+
+	d, err := rpc.NewDRPCServer(
+		ctx,
+		rpcCtx,
+		rpc.WithInterceptor(
+			func(path string) error {
+				return drpcServer.intercept(path)
+			}),
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -41,12 +51,8 @@ func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error
 		return nil, err
 	}
 
-	drpcServer := &drpcServer{
-		DRPCServer: d,
-		tlsCfg:     tlsCfg,
-	}
-
-	drpcServer.setMode(modeInitializing)
+	drpcServer.DRPCServer = d
+	drpcServer.tlsCfg = tlsCfg
 
 	if err := rpc.DRPCRegisterHeartbeat(drpcServer, rpcCtx.NewHeartbeatService()); err != nil {
 		return nil, err

pkg/server/grpc_server.go

@@ -68,24 +68,6 @@ func (s *grpcServer) health(ctx context.Context) error {
 	}
 }
 
-var rpcsAllowedWhileBootstrapping = map[string]struct{}{
-	"/cockroach.rpc.Heartbeat/Ping":             {},
-	"/cockroach.gossip.Gossip/Gossip":           {},
-	"/cockroach.server.serverpb.Init/Bootstrap": {},
-	"/cockroach.server.serverpb.Admin/Health":   {},
-}
-
-// intercept implements filtering rules for each server state.
-func (s *grpcServer) intercept(fullName string) error {
-	if s.operational() {
-		return nil
-	}
-	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
-		return NewWaitingForInitError(fullName)
-	}
-	return nil
-}
-
 // NewWaitingForInitError creates an error indicating that the server cannot run
 // the specified method until the node has been initialized.
 func NewWaitingForInitError(methodName string) error {

pkg/server/serve_mode.go

@@ -36,6 +36,24 @@ const (
 	modeDraining
 )
 
+var rpcsAllowedWhileBootstrapping = map[string]struct{}{
+	"/cockroach.rpc.Heartbeat/Ping":             {},
+	"/cockroach.gossip.Gossip/Gossip":           {},
+	"/cockroach.server.serverpb.Init/Bootstrap": {},
+	"/cockroach.server.serverpb.Admin/Health":   {},
+}
+
+// intercept implements filtering rules for each server state.
+func (s *serveModeHandler) intercept(fullName string) error {
+	if s.operational() {
+		return nil
+	}
+	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
+		return NewWaitingForInitError(fullName)
+	}
+	return nil
+}
+
 func (s *serveMode) set(mode serveMode) {
 	atomic.StoreInt32((*int32)(s), int32(mode))
 }

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/execinfra/base.go

@@ -332,45 +332,6 @@ func DrainAndClose(
 	dst.ProducerDone()
 }
 
-// NoMetadataRowSource is a wrapper on top of a RowSource that automatically
-// forwards metadata to a RowReceiver. Data rows are returned through an
-// interface similar to RowSource, except that, since metadata is taken care of,
-// only the data rows are returned.
-//
-// The point of this struct is that it'd be burdensome for some row consumers to
-// have to deal with metadata.
-type NoMetadataRowSource struct {
-	src          RowSource
-	metadataSink RowReceiver
-}
-
-// MakeNoMetadataRowSource builds a NoMetadataRowSource.
-func MakeNoMetadataRowSource(src RowSource, sink RowReceiver) NoMetadataRowSource {
-	return NoMetadataRowSource{src: src, metadataSink: sink}
-}
-
-// NextRow is analogous to RowSource.Next. If the producer sends an error, we
-// can't just forward it to metadataSink. We need to let the consumer know so
-// that it's not under the impression that everything is hunky-dory and it can
-// continue consuming rows. So, this interface returns the error. Just like with
-// a raw RowSource, the consumer should generally call ConsumerDone() and drain.
-func (rs *NoMetadataRowSource) NextRow() (rowenc.EncDatumRow, error) {
-	for {
-		row, meta := rs.src.Next()
-		if meta == nil {
-			return row, nil
-		}
-		if meta.Err != nil {
-			return nil, meta.Err
-		}
-		// We forward the metadata and ignore the returned ConsumerStatus. There's
-		// no good way to use that status here; eventually the consumer of this
-		// NoMetadataRowSource will figure out the same status and act on it as soon
-		// as a non-metadata row is received.
-		_ = rs.metadataSink.Push(nil /* row */, meta)
-	}
-}
-
 // RowChannelMsg is the message used in the channels that implement
 // local physical streams (i.e. the RowChannel's).
 type RowChannelMsg struct {

pkg/sql/export/BUILD.bazel

@@ -15,16 +15,17 @@ go_library(
         "//pkg/settings",
         "//pkg/sql/catalog/colinfo",
         "//pkg/sql/execinfra",
+        "//pkg/sql/execinfra/execopnode",
         "//pkg/sql/execinfrapb",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
         "//pkg/util/encoding/csv",
+        "//pkg/util/log",
         "//pkg/util/mon",
         "//pkg/util/parquet",
-        "//pkg/util/tracing",
         "//pkg/util/unique",
         "@com_github_cockroachdb_errors//:errors",
     ],