colexec: add more redundancy to releasing disk resources

pkg/sql/colexec/colbuilder/execplan.go

@@ -943,7 +943,7 @@ func NewColOperator(
 							newAggArgs.Allocator = colmem.NewAllocator(ctx, ehaMemAccount, factory)
 							newAggArgs.MemAccount = ehaMemAccount
 							newAggArgs.Input = input
-							return colexecdisk.NewExternalHashAggregator(
+							eha, toClose := colexecdisk.NewExternalHashAggregator(
 								flowCtx,
 								args,
 								&newAggArgs,
@@ -957,6 +957,8 @@ func NewColOperator(
 								outputUnlimitedAllocator,
 								maxOutputBatchMemSize,
 							)
+							result.ToClose = append(result.ToClose, toClose)
+							return eha
 						},
 						args.TestingKnobs.SpillingCallbackFn,
 					)
@@ -1006,7 +1008,7 @@ func NewColOperator(
 						unlimitedAllocator := colmem.NewAllocator(
 							ctx, args.MonitorRegistry.CreateUnlimitedMemAccount(ctx, flowCtx, edOpName, spec.ProcessorID), factory,
 						)
-						return colexecdisk.NewExternalDistinct(
+						ed, toClose := colexecdisk.NewExternalDistinct(
 							unlimitedAllocator,
 							flowCtx,
 							args,
@@ -1016,6 +1018,8 @@ func NewColOperator(
 							inMemoryUnorderedDistinct,
 							diskAccount,
 						)
+						result.ToClose = append(result.ToClose, toClose)
+						return ed
 					},
 					args.TestingKnobs.SpillingCallbackFn,
 				)
@@ -1106,7 +1110,7 @@ func NewColOperator(
 								result.makeDiskBackedSorterConstructor(ctx, flowCtx, args, opName, factory),
 								diskAccount,
 							)
-							result.ToClose = append(result.ToClose, ehj.(colexecop.Closer))
+							result.ToClose = append(result.ToClose, ehj)
 							return ehj
 						},
 						args.TestingKnobs.SpillingCallbackFn,
@@ -1419,7 +1423,9 @@ func NewColOperator(
 						}
 						result.Root = colexecwindow.NewWindowAggregatorOperator(
 							windowArgs, aggType, wf.Frame, &wf.Ordering, argIdxs,
-							aggArgs.OutputTypes[0], aggFnsAlloc, toClose)
+							aggArgs.OutputTypes[0], aggFnsAlloc,
+						)
+						result.ToClose = append(result.ToClose, toClose...)
 						returnType = aggArgs.OutputTypes[0]
 					}
 				} else {

pkg/sql/colexec/colexecdisk/external_distinct.go

@@ -36,7 +36,7 @@ func NewExternalDistinct(
 	createDiskBackedSorter DiskBackedSorterConstructor,
 	inMemUnorderedDistinct colexecop.Operator,
 	diskAcc *mon.BoundAccount,
-) colexecop.Operator {
+) (colexecop.Operator, colexecop.Closer) {
 	distinctSpec := args.Spec.Core.Distinct
 	distinctCols := distinctSpec.DistinctColumns
 	inMemMainOpConstructor := func(partitionedInputs []*partitionerToOperator) colexecop.ResettableOperator {
@@ -111,13 +111,13 @@ func NewExternalDistinct(
 	outputOrdering := args.Spec.Core.Distinct.OutputOrdering
 	if len(outputOrdering.Columns) == 0 {
 		// No particular output ordering is required.
-		return ed
+		return ed, ed
 	}
 	// TODO(yuzefovich): the fact that we're planning an additional external
 	// sort isn't accounted for when considering the number file descriptors to
 	// acquire. Not urgent, but it should be fixed.
 	maxNumberActivePartitions := calculateMaxNumberActivePartitions(flowCtx, args, numRequiredActivePartitions)
-	return createDiskBackedSorter(ed, inputTypes, outputOrdering.Columns, maxNumberActivePartitions)
+	return createDiskBackedSorter(ed, inputTypes, outputOrdering.Columns, maxNumberActivePartitions), ed
 }
 
 // unorderedDistinctFilterer filters out tuples that are duplicates of the

pkg/sql/colexec/colexecdisk/external_hash_aggregator.go

@@ -39,7 +39,7 @@ func NewExternalHashAggregator(
 	diskAcc *mon.BoundAccount,
 	outputUnlimitedAllocator *colmem.Allocator,
 	maxOutputBatchMemSize int64,
-) colexecop.Operator {
+) (colexecop.Operator, colexecop.Closer) {
 	inMemMainOpConstructor := func(partitionedInputs []*partitionerToOperator) colexecop.ResettableOperator {
 		newAggArgs := *newAggArgs
 		newAggArgs.Input = partitionedInputs[0]
@@ -83,11 +83,11 @@ func NewExternalHashAggregator(
 	outputOrdering := args.Spec.Core.Aggregator.OutputOrdering
 	if len(outputOrdering.Columns) == 0 {
 		// No particular output ordering is required.
-		return eha
+		return eha, eha
 	}
 	// TODO(yuzefovich): the fact that we're planning an additional external
 	// sort isn't accounted for when considering the number file descriptors to
 	// acquire. Not urgent, but it should be fixed.
 	maxNumberActivePartitions := calculateMaxNumberActivePartitions(flowCtx, args, ehaNumRequiredActivePartitions)
-	return createDiskBackedSorter(eha, newAggArgs.OutputTypes, outputOrdering.Columns, maxNumberActivePartitions)
+	return createDiskBackedSorter(eha, newAggArgs.OutputTypes, outputOrdering.Columns, maxNumberActivePartitions), eha
 }

pkg/sql/colexec/colexecdisk/external_hash_joiner.go

@@ -68,7 +68,7 @@ func NewExternalHashJoiner(
 	leftInput, rightInput colexecop.Operator,
 	createDiskBackedSorter DiskBackedSorterConstructor,
 	diskAcc *mon.BoundAccount,
-) colexecop.Operator {
+) colexecop.ClosableOperator {
 	// This memory limit will restrict the size of the batches output by the
 	// in-memory hash joiner in the main strategy as well as by the merge joiner
 	// in the fallback strategy.

pkg/sql/colexec/colexecjoin/crossjoiner.go

@@ -39,7 +39,7 @@ func NewCrossJoiner(
 	leftTypes []*types.T,
 	rightTypes []*types.T,
 	diskAcc *mon.BoundAccount,
-) colexecop.Operator {
+) colexecop.ClosableOperator {
 	return &crossJoiner{
 		crossJoinerBase: newCrossJoinerBase(
 			unlimitedAllocator,

pkg/sql/colexec/colexecwindow/buffered_window.go

@@ -30,7 +30,7 @@ import (
 // window function.
 func newBufferedWindowOperator(
 	args *WindowArgs, windower bufferedWindower, outputColType *types.T, memoryLimit int64,
-) colexecop.Operator {
+) colexecop.ClosableOperator {
 	outputTypes := make([]*types.T, len(args.InputTypes), len(args.InputTypes)+1)
 	copy(outputTypes, args.InputTypes)
 	outputTypes = append(outputTypes, outputColType)
@@ -213,7 +213,7 @@ func (b *bufferedWindowOp) Init(ctx context.Context) {
 	b.windower.startNewPartition()
 }
 
-var _ colexecop.Operator = &bufferedWindowOp{}
+var _ colexecop.ClosableOperator = &bufferedWindowOp{}
 
 func (b *bufferedWindowOp) Next() coldata.Batch {
 	var err error
@@ -343,16 +343,11 @@ func (b *bufferedWindowOp) Next() coldata.Batch {
 }
 
 func (b *bufferedWindowOp) Close(ctx context.Context) error {
-	if !b.CloserHelper.Close() || b.Ctx == nil {
-		// Either Close() has already been called or Init() was never called. In
-		// both cases there is nothing to do.
+	if !b.CloserHelper.Close() {
 		return nil
 	}
-	if err := b.bufferQueue.Close(ctx); err != nil {
-		return err
-	}
 	b.windower.Close(ctx)
-	return nil
+	return b.bufferQueue.Close(ctx)
 }
 
 // partitionSeekerBase extracts common fields and methods for buffered windower

pkg/sql/colexec/colexecwindow/count_rows_aggregator.go

@@ -25,7 +25,7 @@ import (
 // aggregate window function.
 func NewCountRowsOperator(
 	args *WindowArgs, frame *execinfrapb.WindowerSpec_Frame, ordering *execinfrapb.Ordering,
-) colexecop.Operator {
+) colexecop.ClosableOperator {
 	// Because the buffer is potentially used multiple times per-row, it is
 	// important to prevent it from spilling to disk if possible. For this reason,
 	// we give the buffer half of the memory budget even though it will generally

pkg/sql/colexec/colexecwindow/first_last_nth_value_tmpl.go

@@ -56,7 +56,7 @@ func New_UPPERCASE_NAMEOperator(
 	frame *execinfrapb.WindowerSpec_Frame,
 	ordering *execinfrapb.Ordering,
 	argIdxs []int,
-) (colexecop.Operator, error) {
+) (colexecop.ClosableOperator, error) {
 	framer := newWindowFramer(args.EvalCtx, frame, ordering, args.InputTypes, args.PeersColIdx)
 	colsToStore := framer.getColsToStore([]int{argIdxs[0]})
 

pkg/sql/colexec/colexecwindow/lead_lag_tmpl.go

@@ -49,7 +49,7 @@ const _TYPE_WIDTH = 0
 // should put its output (if there is no such column, a new column is appended).
 func New_UPPERCASE_NAMEOperator(
 	args *WindowArgs, argIdx int, offsetIdx int, defaultIdx int,
-) (colexecop.Operator, error) {
+) (colexecop.ClosableOperator, error) {
 	// Allow the direct-access buffer 10% of the available memory. The rest will
 	// be given to the bufferedWindowOp queue. While it is somewhat more important
 	// for the direct-access buffer tuples to be kept in-memory, it only has to

pkg/sql/colexec/colexecwindow/relative_rank_tmpl.go

@@ -590,9 +590,7 @@ func (r *_RELATIVE_RANK_STRINGOp) Next() coldata.Batch {
 }
 
 func (r *_RELATIVE_RANK_STRINGOp) Close(ctx context.Context) error {
-	if !r.CloserHelper.Close() || r.Ctx == nil {
-		// Either Close() has already been called or Init() was never called. In
-		// both cases there is nothing to do.
+	if !r.CloserHelper.Close() {
 		return nil
 	}
 	var lastErr error

pkg/sql/colexec/colexecwindow/window_aggregator_tmpl.go

@@ -27,7 +27,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecagg"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecutils"
-	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
 	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
@@ -57,8 +56,7 @@ func NewWindowAggregatorOperator(
 	argIdxs []int,
 	outputType *types.T,
 	aggAlloc *colexecagg.AggregateFuncsAlloc,
-	closers colexecop.Closers,
-) colexecop.Operator {
+) colexecop.ClosableOperator {
 	// Because the buffer is used multiple times per-row, it is important to
 	// prevent it from spilling to disk if possible. For this reason, we give the
 	// buffer half of the memory budget even though it will generally store less
@@ -84,7 +82,6 @@ func NewWindowAggregatorOperator(
 		outputColIdx: args.OutputColIdx,
 		inputIdxs:    inputIdxs,
 		framer:       framer,
-		closers:      closers,
 		vecs:         make([]coldata.Vec, len(inputIdxs)),
 	}
 	var agg colexecagg.AggregateFunc
@@ -131,7 +128,6 @@ func NewWindowAggregatorOperator(
 type windowAggregatorBase struct {
 	partitionSeekerBase
 	colexecop.CloserHelper
-	closers   colexecop.Closers
 	allocator *colmem.Allocator
 
 	outputColIdx int
@@ -183,9 +179,6 @@ func (a *windowAggregatorBase) Close(ctx context.Context) {
 	if !a.CloserHelper.Close() {
 		return
 	}
-	if err := a.closers.Close(ctx); err != nil {
-		colexecerror.InternalError(err)
-	}
 	a.framer.close()
 	a.buffer.Close(ctx)
 }

pkg/sql/colexec/colexecwindow/window_functions_test.go

@@ -1124,6 +1124,11 @@ func BenchmarkWindowFunctions(b *testing.B) {
 	benchMemAccount := testMemMonitor.MakeBoundAccount()
 	defer benchMemAccount.Close(ctx)
 
+	var allClosers colexecop.Closers
+	defer func() {
+		require.NoError(b, allClosers.Close(ctx))
+	}()
+
 	getWindowFn := func(
 		fun execinfrapb.WindowerSpec_Func, source colexecop.Operator, partition, order bool,
 	) (op colexecop.Operator) {
@@ -1222,7 +1227,9 @@ func BenchmarkWindowFunctions(b *testing.B) {
 			op = NewWindowAggregatorOperator(
 				args, *fun.AggregateFunc, NormalizeWindowFrame(nil),
 				&execinfrapb.Ordering{Columns: orderingCols}, []int{arg1ColIdx},
-				aggArgs.OutputTypes[0], aggFnsAlloc, toClose)
+				aggArgs.OutputTypes[0], aggFnsAlloc,
+			)
+			allClosers = append(allClosers, toClose...)
 		} else {
 			require.Fail(b, "expected non-nil window function")
 		}