sql,server: add admission control for DistSQL responses

The admission.WorkQueue is plumbed through via execinfra.ServerCfg
to colrpc.Inbox and flowinfra.processProducerMessage, and the
AdmissionInfo is initialized using FlowCtx.Txn when possible.

The admission control point for KV=>SQL response processing
is tweaked to move it after the memory accounting for the
response bytes.

Informs cockroachdb#65948

Release note (ops change): DistSQL response admission control
can be enabled using admission.sql_sql_response.enabled.

docs/generated/settings/settings-for-tenants.txt

@@ -1,6 +1,7 @@
 Setting	Type	Default	Description
 admission.kv.enabled	boolean	false	when true, work performed by the KV layer is subject to admission control
 admission.sql_kv_response.enabled	boolean	false	when true, work performed by the SQL layer when receiving a KV response is subject to admission control
+admission.sql_sql_response.enabled	boolean	false	when true, work performed by the SQL layer when receiving a DistSQL response is subject to admission control
 bulkio.stream_ingestion.minimum_flush_interval	duration	5s	the minimum timestamp between flushes; flushes may still occur if internal buffers fill up
 cloudstorage.http.custom_ca	string		custom root CA (appended to system's default CAs) for verifying certificates when interacting with HTTPS storage
 cloudstorage.timeout	duration	10m0s	the timeout for import/export storage operations

docs/generated/settings/settings.html

@@ -3,6 +3,7 @@
 <tbody>
 <tr><td><code>admission.kv.enabled</code></td><td>boolean</td><td><code>false</code></td><td>when true, work performed by the KV layer is subject to admission control</td></tr>
 <tr><td><code>admission.sql_kv_response.enabled</code></td><td>boolean</td><td><code>false</code></td><td>when true, work performed by the SQL layer when receiving a KV response is subject to admission control</td></tr>
+<tr><td><code>admission.sql_sql_response.enabled</code></td><td>boolean</td><td><code>false</code></td><td>when true, work performed by the SQL layer when receiving a DistSQL response is subject to admission control</td></tr>
 <tr><td><code>bulkio.stream_ingestion.minimum_flush_interval</code></td><td>duration</td><td><code>5s</code></td><td>the minimum timestamp between flushes; flushes may still occur if internal buffers fill up</td></tr>
 <tr><td><code>cloudstorage.http.custom_ca</code></td><td>string</td><td><code></code></td><td>custom root CA (appended to system's default CAs) for verifying certificates when interacting with HTTPS storage</td></tr>
 <tr><td><code>cloudstorage.timeout</code></td><td>duration</td><td><code>10m0s</code></td><td>the timeout for import/export storage operations</td></tr>

pkg/server/server.go

@@ -706,14 +706,15 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 
 	sqlServer, err := newSQLServer(ctx, sqlServerArgs{
 		sqlServerOptionalKVArgs: sqlServerOptionalKVArgs{
-			nodesStatusServer:      serverpb.MakeOptionalNodesStatusServer(sStatus),
-			nodeLiveness:           optionalnodeliveness.MakeContainer(nodeLiveness),
-			gossip:                 gossip.MakeOptionalGossip(g),
-			grpcServer:             grpcServer.Server,
-			nodeIDContainer:        idContainer,
-			externalStorage:        externalStorage,
-			externalStorageFromURI: externalStorageFromURI,
-			isMeta1Leaseholder:     node.stores.IsMeta1Leaseholder,
+			nodesStatusServer:        serverpb.MakeOptionalNodesStatusServer(sStatus),
+			nodeLiveness:             optionalnodeliveness.MakeContainer(nodeLiveness),
+			gossip:                   gossip.MakeOptionalGossip(g),
+			grpcServer:               grpcServer.Server,
+			nodeIDContainer:          idContainer,
+			externalStorage:          externalStorage,
+			externalStorageFromURI:   externalStorageFromURI,
+			isMeta1Leaseholder:       node.stores.IsMeta1Leaseholder,
+			sqlSQLResponseAdmissionQ: gcoord.GetWorkQueue(admission.SQLSQLResponseWork),
 		},
 		SQLConfig:                &cfg.SQLConfig,
 		BaseConfig:               &cfg.BaseConfig,

pkg/server/server_sql.go

@@ -74,6 +74,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sqlmigrations"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/cloud"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -164,6 +165,9 @@ type sqlServerOptionalKVArgs struct {
 	// Used by backup/restore.
 	externalStorage        cloud.ExternalStorageFactory
 	externalStorageFromURI cloud.ExternalStorageFromURIFactory
+
+	// The admission queue to use for SQLSQLResponseWork.
+	sqlSQLResponseAdmissionQ *admission.WorkQueue
 }
 
 // sqlServerOptionalTenantArgs are the arguments supplied to newSQLServer which
@@ -469,9 +473,10 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		ExternalStorage:        cfg.externalStorage,
 		ExternalStorageFromURI: cfg.externalStorageFromURI,
 
-		RangeCache:     cfg.distSender.RangeDescriptorCache(),
-		HydratedTables: hydratedTablesCache,
-		VirtualSchemas: virtualSchemas,
+		RangeCache:               cfg.distSender.RangeDescriptorCache(),
+		HydratedTables:           hydratedTablesCache,
+		VirtualSchemas:           virtualSchemas,
+		SQLSQLResponseAdmissionQ: cfg.sqlSQLResponseAdmissionQ,
 	}
 	cfg.TempStorageConfig.Mon.SetMetrics(distSQLMetrics.CurDiskBytesCount, distSQLMetrics.MaxDiskBytesHist)
 	if distSQLTestingKnobs := cfg.TestingKnobs.DistSQL; distSQLTestingKnobs != nil {

pkg/sql/colflow/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/sql/sessiondatapb",
         "//pkg/sql/types",
         "//pkg/util",
+        "//pkg/util/admission",
         "//pkg/util/log",
         "//pkg/util/metric",
         "//pkg/util/mon",
@@ -108,6 +109,7 @@ go_test(
         "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
+        "//pkg/util/admission",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
         "//pkg/util/leaktest",

pkg/sql/colflow/colrpc/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/types",
+        "//pkg/util/admission",
         "//pkg/util/cancelchecker",
         "//pkg/util/log",
         "//pkg/util/log/logcrash",

pkg/sql/colflow/colrpc/inbox.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logcrash"
@@ -98,6 +99,9 @@ type Inbox struct {
 	// only the Next/DrainMeta goroutine may access it.
 	stream flowStreamServer
 
+	admissionQ    *admission.WorkQueue
+	admissionInfo admission.WorkInfo
+
 	// statsAtomics are the execution statistics that need to be atomically
 	// accessed. This is necessary since Get*() methods can be called from
 	// different goroutine than Next().
@@ -151,6 +155,24 @@ func NewInbox(
 	return i, nil
 }
 
+// NewInboxWithAdmissionControl creates a new Inbox that does admission
+// control on responses received from DistSQL.
+func NewInboxWithAdmissionControl(
+	allocator *colmem.Allocator,
+	typs []*types.T,
+	streamID execinfrapb.StreamID,
+	admissionQ *admission.WorkQueue,
+	admissionInfo admission.WorkInfo,
+) (*Inbox, error) {
+	i, err := NewInbox(allocator, typs, streamID)
+	if err != nil {
+		return nil, err
+	}
+	i.admissionQ = admissionQ
+	i.admissionInfo = admissionInfo
+	return i, err
+}
+
 // close closes the inbox, ensuring that any call to RunWithStream will return
 // immediately. close is idempotent.
 // NOTE: it is very important to close the Inbox only when execution terminates
@@ -333,6 +355,15 @@ func (i *Inbox) Next() coldata.Batch {
 		// Update the allocator since we're holding onto the serialized bytes
 		// for now.
 		i.allocator.AdjustMemoryUsage(numSerializedBytes)
+		// Do admission control after memory accounting for the serialized bytes
+		// and before deserialization.
+		if i.admissionQ != nil {
+			if _, err := i.admissionQ.Admit(i.Ctx, i.admissionInfo); err != nil {
+				// err includes the case of context cancellation while waiting for
+				// admission.
+				colexecerror.ExpectedError(err)
+			}
+		}
 		i.scratch.data = i.scratch.data[:0]
 		batchLength, err := i.serializer.Deserialize(&i.scratch.data, m.Data.RawBytes)
 		// Eagerly throw away the RawBytes memory.

pkg/sql/colflow/explain_vec.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/flowinfra"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/treeprinter"
 	"github.com/cockroachdb/errors"
@@ -57,6 +58,7 @@ func convertToVecTree(
 		newNoopFlowCreatorHelper(), vectorizedRemoteComponentCreator{}, false, false,
 		nil, &execinfra.RowChannel{}, &fakeBatchReceiver{}, nil, execinfrapb.FlowID{}, colcontainer.DiskQueueCfg{},
 		flowCtx.Cfg.VecFDSemaphore, flowCtx.TypeResolverFactory.NewTypeResolver(flowCtx.EvalCtx.Txn),
+		admission.WorkInfo{},
 	)
 	// We create an unlimited memory account because we're interested whether the
 	// flow is supported via the vectorized engine in general (without paying

pkg/sql/colflow/vectorized_flow.go

@@ -40,6 +40,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
@@ -208,6 +209,7 @@ func (f *vectorizedFlow) Setup(
 		diskQueueCfg,
 		f.countingSemaphore,
 		flowCtx.TypeResolverFactory.NewTypeResolver(flowCtx.EvalCtx.Txn),
+		f.FlowBase.GetAdmissionInfo(),
 	)
 	if f.testingKnobs.onSetupFlow != nil {
 		f.testingKnobs.onSetupFlow(f.creator)
@@ -449,6 +451,11 @@ type flowCreatorHelper interface {
 	getCancelFlowFn() context.CancelFunc
 }
 
+type admissionOptions struct {
+	admissionQ    *admission.WorkQueue
+	admissionInfo admission.WorkInfo
+}
+
 // remoteComponentCreator is an interface that abstracts the constructors for
 // several components in a remote flow. Mostly for testing purposes.
 type remoteComponentCreator interface {
@@ -460,7 +467,8 @@ type remoteComponentCreator interface {
 		metadataSources []colexecop.MetadataSource,
 		toClose []colexecop.Closer,
 	) (*colrpc.Outbox, error)
-	newInbox(allocator *colmem.Allocator, typs []*types.T, streamID execinfrapb.StreamID) (*colrpc.Inbox, error)
+	newInbox(allocator *colmem.Allocator, typs []*types.T, streamID execinfrapb.StreamID,
+		admissionOpts admissionOptions) (*colrpc.Inbox, error)
 }
 
 type vectorizedRemoteComponentCreator struct{}
@@ -477,9 +485,13 @@ func (vectorizedRemoteComponentCreator) newOutbox(
 }
 
 func (vectorizedRemoteComponentCreator) newInbox(
-	allocator *colmem.Allocator, typs []*types.T, streamID execinfrapb.StreamID,
+	allocator *colmem.Allocator,
+	typs []*types.T,
+	streamID execinfrapb.StreamID,
+	admissionOpts admissionOptions,
 ) (*colrpc.Inbox, error) {
-	return colrpc.NewInbox(allocator, typs, streamID)
+	return colrpc.NewInboxWithAdmissionControl(
+		allocator, typs, streamID, admissionOpts.admissionQ, admissionOpts.admissionInfo)
 }
 
 // vectorizedFlowCreator performs all the setup of vectorized flows. Depending
@@ -508,6 +520,7 @@ type vectorizedFlowCreator struct {
 	flowID            execinfrapb.FlowID
 	exprHelper        *colexecargs.ExprHelper
 	typeResolver      descs.DistSQLTypeResolver
+	admissionInfo     admission.WorkInfo
 
 	// numOutboxes counts how many colrpc.Outbox'es have been set up on this
 	// node. It must be accessed atomically.
@@ -573,6 +586,7 @@ func newVectorizedFlowCreator(
 	diskQueueCfg colcontainer.DiskQueueCfg,
 	fdSemaphore semaphore.Semaphore,
 	typeResolver descs.DistSQLTypeResolver,
+	admissionInfo admission.WorkInfo,
 ) *vectorizedFlowCreator {
 	creator := vectorizedFlowCreatorPool.Get().(*vectorizedFlowCreator)
 	*creator = vectorizedFlowCreator{
@@ -589,6 +603,7 @@ func newVectorizedFlowCreator(
 		flowID:                 flowID,
 		exprHelper:             creator.exprHelper,
 		typeResolver:           typeResolver,
+		admissionInfo:          admissionInfo,
 		procIdxQueue:           creator.procIdxQueue,
 		opChains:               creator.opChains,
 		monitors:               creator.monitors,
@@ -881,8 +896,14 @@ func (s *vectorizedFlowCreator) setupInput(
 				latency = 0
 				log.VEventf(ctx, 1, "an error occurred during vectorized planning while getting latency: %v", err)
 			}
-
-			inbox, err := s.remoteComponentCreator.newInbox(colmem.NewAllocator(ctx, s.newStreamingMemAccount(flowCtx), factory), input.ColumnTypes, inputStream.StreamID)
+			inbox, err := s.remoteComponentCreator.newInbox(
+				colmem.NewAllocator(ctx, s.newStreamingMemAccount(flowCtx), factory),
+				input.ColumnTypes,
+				inputStream.StreamID,
+				admissionOptions{
+					admissionQ:    flowCtx.Cfg.SQLSQLResponseAdmissionQ,
+					admissionInfo: s.admissionInfo,
+				})
 
 			if err != nil {
 				return colexecargs.OpWithMetaInfo{}, err

pkg/sql/colflow/vectorized_flow_test.go

@@ -31,6 +31,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/stretchr/testify/require"
 )
@@ -52,7 +53,7 @@ func (c callbackRemoteComponentCreator) newOutbox(
 }
 
 func (c callbackRemoteComponentCreator) newInbox(
-	allocator *colmem.Allocator, typs []*types.T, streamID execinfrapb.StreamID,
+	allocator *colmem.Allocator, typs []*types.T, streamID execinfrapb.StreamID, _ admissionOptions,
 ) (*colrpc.Inbox, error) {
 	return c.newInboxFn(allocator, typs, streamID)
 }
@@ -218,15 +219,17 @@ func TestDrainOnlyInputDAG(t *testing.T) {
 	ctx := context.Background()
 	defer evalCtx.Stop(ctx)
 	f := &flowinfra.FlowBase{
-		FlowCtx: execinfra.FlowCtx{EvalCtx: &evalCtx,
-			NodeID: base.TestingIDContainer,
+		FlowCtx: execinfra.FlowCtx{
+			Cfg:     &execinfra.ServerConfig{},
+			EvalCtx: &evalCtx,
+			NodeID:  base.TestingIDContainer,
 		},
 	}
 	var wg sync.WaitGroup
 	vfc := newVectorizedFlowCreator(
 		&vectorizedFlowCreatorHelper{f: f}, componentCreator, false, false, &wg, &execinfra.RowChannel{},
 		nil /* batchSyncFlowConsumer */, nil /* nodeDialer */, execinfrapb.FlowID{}, colcontainer.DiskQueueCfg{},
-		nil /* fdSemaphore */, descs.DistSQLTypeResolver{},
+		nil /* fdSemaphore */, descs.DistSQLTypeResolver{}, admission.WorkInfo{},
 	)
 
 	_, _, err := vfc.setupFlow(ctx, &f.FlowCtx, procs, nil /* localProcessors */, flowinfra.FuseNormally)

pkg/sql/execinfra/BUILD.bazel

@@ -52,6 +52,7 @@ go_library(
         "//pkg/storage/cloud",
         "//pkg/storage/fs",
         "//pkg/util",
+        "//pkg/util/admission",
         "//pkg/util/log",
         "//pkg/util/log/logcrash",
         "//pkg/util/metric",

pkg/sql/execinfra/server_config.go

@@ -35,6 +35,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
 	"github.com/cockroachdb/cockroach/pkg/storage/cloud"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
@@ -151,6 +152,10 @@ type ServerConfig struct {
 
 	// VirtualSchemas hold the virtual table schemas.
 	VirtualSchemas catalog.VirtualSchemas
+
+	// SQLSQLResponseAdmissionQ is the admission queue to use for
+	// SQLSQLResponseWork.
+	SQLSQLResponseAdmissionQ *admission.WorkQueue
 }
 
 // RuntimeStats is an interface through which the rowexec layer can get

pkg/sql/flowinfra/BUILD.bazel

@@ -27,6 +27,7 @@ go_library(
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/rowenc",
         "//pkg/sql/types",
+        "//pkg/util/admission",
         "//pkg/util/cancelchecker",
         "//pkg/util/contextutil",
         "//pkg/util/log",