Merge #67531

67531: sql,server: add admission control for DistSQL responses r=sumeerbhola a=sumeerbhola

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 #65948

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

Co-authored-by: sumeerbhola <[email protected]>

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/startupmigrations"
 	"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

@@ -53,6 +53,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",