Merge #26914

26914: distsql: add memory accounting to mergeJoiner r=yuzefovich a=yuzefovich

Resolves: #26896.

Adds accounting of memory allocated during computation of
cartesion product of equal rows and adds the max allocated
memory to the visual query plan stats.

Release note: None

Co-authored-by: yuzefovich <yahor@cockroachlabs.com>

Author: craig[bot]

pkg/sql/distsqlrun/mergejoiner.go

@@ -19,8 +19,11 @@ import (
 	"errors"
 	"sync"
 
+	"fmt"
+
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/opentracing/opentracing-go"
 )
@@ -84,18 +87,25 @@ func newMergeJoiner(
 		spec.Type, spec.OnExpr, leftEqCols, rightEqCols, 0, post, output,
 		procStateOpts{
 			inputsToDrain: []RowSource{leftSource, rightSource},
+			trailingMetaCallback: func() []ProducerMetadata {
+				m.close()
+				return nil
+			},
 		},
 	); err != nil {
 		return nil, err
 	}
 
+	m.memMonitor = newMonitor(flowCtx.EvalCtx.Ctx(), flowCtx.EvalCtx.Mon, "mergejoiner-mem")
+
 	var err error
 	m.streamMerger, err = makeStreamMerger(
 		m.leftSource,
 		convertToColumnOrdering(spec.LeftOrdering),
 		m.rightSource,
 		convertToColumnOrdering(spec.RightOrdering),
 		spec.NullEquality,
+		m.memMonitor,
 	)
 	if err != nil {
 		return nil, err
@@ -248,10 +258,18 @@ func (m *mergeJoiner) ConsumerDone() {
 	m.moveToDraining(nil /* err */)
 }
 
+func (m *mergeJoiner) close() {
+	if m.internalClose() {
+		ctx := m.evalCtx.Ctx()
+		m.streamMerger.close(ctx)
+		m.memMonitor.Stop(ctx)
+	}
+}
+
 // ConsumerClosed is part of the RowSource interface.
 func (m *mergeJoiner) ConsumerClosed() {
 	// The consumer is done, Next() will not be called again.
-	m.internalClose()
+	m.close()
 }
 
 var _ DistSQLSpanStats = &MergeJoinerStats{}
@@ -267,15 +285,17 @@ func (mjs *MergeJoinerStats) Stats() map[string]string {
 	for k, v := range rightInputStatsMap {
 		statsMap[k] = v
 	}
+	statsMap[mergeJoinerTagPrefix+maxMemoryTagSuffix] = humanizeutil.IBytes(mjs.MaxAllocatedMem)
 	return statsMap
 }
 
 // StatsForQueryPlan implements the DistSQLSpanStats interface.
 func (mjs *MergeJoinerStats) StatsForQueryPlan() []string {
-	return append(
+	stats := append(
 		mjs.LeftInputStats.StatsForQueryPlan("left "),
 		mjs.RightInputStats.StatsForQueryPlan("right ")...,
 	)
+	return append(stats, fmt.Sprintf("%s: %s", maxMemoryQueryPlanSuffix, humanizeutil.IBytes(mjs.MaxAllocatedMem)))
 }
 
 // outputStatsToTrace outputs the collected mergeJoiner stats to the trace. Will
@@ -295,6 +315,7 @@ func (m *mergeJoiner) outputStatsToTrace() {
 			&MergeJoinerStats{
 				LeftInputStats:  lis,
 				RightInputStats: ris,
+				MaxAllocatedMem: m.memMonitor.MaximumBytes(),
 			},
 		)
 	}

pkg/sql/distsqlrun/stats.pb.go

@@ -58,6 +58,7 @@ message DistinctStats {
 message MergeJoinerStats {
   InputStats left_input_stats = 1 [(gogoproto.nullable) = false];
   InputStats right_input_stats = 2 [(gogoproto.nullable) = false];
+  int64 max_allocated_mem = 3;
 }
 
 // SorterStats are the stats collected during a sorter run.

pkg/sql/distsqlrun/stats.proto

@@ -19,6 +19,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
+	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/pkg/errors"
 )
 
@@ -42,15 +43,18 @@ type streamGroupAccumulator struct {
 	leftoverRow sqlbase.EncDatumRow
 
 	rowAlloc sqlbase.EncDatumRowAlloc
+
+	memAcc mon.BoundAccount
 }
 
 func makeStreamGroupAccumulator(
-	src RowSource, ordering sqlbase.ColumnOrdering,
+	src RowSource, ordering sqlbase.ColumnOrdering, memMonitor *mon.BytesMonitor,
 ) streamGroupAccumulator {
 	return streamGroupAccumulator{
 		src:      src,
 		types:    src.OutputTypes(),
 		ordering: ordering,
+		memAcc:   memMonitor.MakeBoundAccount(),
 	}
 }
 
@@ -84,6 +88,9 @@ func (s *streamGroupAccumulator) nextGroup(
 			return s.curGroup, nil
 		}
 
+		if err := s.memAcc.Grow(evalCtx.Ctx(), int64(row.Size())); err != nil {
+			return nil, &ProducerMetadata{Err: err}
+		}
 		row = s.rowAlloc.CopyRow(row)
 
 		if len(s.curGroup) == 0 {
@@ -115,3 +122,7 @@ func (s *streamGroupAccumulator) nextGroup(
 		}
 	}
 }
+
+func (s *streamGroupAccumulator) close(ctx context.Context) {
+	s.memAcc.Close(ctx)
+}

pkg/sql/distsqlrun/stream_group_accumulator.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
+	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/pkg/errors"
 )
 
@@ -156,6 +157,11 @@ func CompareEncDatumRowForMerge(
 	return 0, nil
 }
 
+func (sm *streamMerger) close(ctx context.Context) {
+	sm.left.close(ctx)
+	sm.right.close(ctx)
+}
+
 // makeStreamMerger creates a streamMerger, joining rows from leftSource with
 // rows from rightSource.
 //
@@ -166,6 +172,7 @@ func makeStreamMerger(
 	rightSource RowSource,
 	rightOrdering sqlbase.ColumnOrdering,
 	nullEquality bool,
+	memMonitor *mon.BytesMonitor,
 ) (streamMerger, error) {
 	if len(leftOrdering) != len(rightOrdering) {
 		return streamMerger{}, errors.Errorf(
@@ -178,8 +185,8 @@ func makeStreamMerger(
 	}
 
 	return streamMerger{
-		left:         makeStreamGroupAccumulator(leftSource, leftOrdering),
-		right:        makeStreamGroupAccumulator(rightSource, rightOrdering),
+		left:         makeStreamGroupAccumulator(leftSource, leftOrdering, memMonitor),
+		right:        makeStreamGroupAccumulator(rightSource, rightOrdering, memMonitor),
 		nullEquality: nullEquality,
 	}, nil
 }

pkg/sql/distsqlrun/stream_merger.go

@@ -139,7 +139,7 @@ INSERT INTO kw VALUES (1,1);
 query T
 SELECT "URL" FROM [EXPLAIN ANALYZE (DISTSQL) SELECT kv.k, avg(kw.k) FROM kv JOIN kw ON kv.k=kw.k GROUP BY kv.k]
 ----
-https://cockroachdb.github.io/distsqlplan/decode.html?eJysk0Fv1DAQhe_8CmtOW2Fp46QnS0hZLqhQElQVJISiysRDsEjiaOzQVtX-dxT7sCRsFhb1tn6eb_zezOYJequxUB06kF9AQMVhIFujc5YmKRZc6QeQCQfTD6Of5IpDbQlBPoE3vkWQcKu-tniDSiNtE-Cg0SvThrYDmU7RY_7jJ3Age-8YodKSCeDgvGpb5k2HkiUOOJSjlywXUO052NEfHnReNQhS7Pn_mRIrpu6fwVS6aurgZewtaSTUMx_VRP6t5Eiy90gNvrWmR9qm82QtfvObXLy8eEWm-R5_Ag8yWwQN2h9pA7YsjeLaZHierg4nO2dju6YhbJS3tM3msfLJRRnnE5bBYVd8vivK27vi4_X1JhdTyt2nN5s8vfiHnXbqgXXYWXpko5taioS9M69XY1yeE-MG3WB7h8tdH-2cTAtG3WD8wzg7Uo0fyNbhmXgsAxcEjc7H2zQervpwFb6M32FxBpwu4fQknM3gZAlnJ-HLBVztX_wKAAD__5wEeH8=
+https://cockroachdb.github.io/distsqlplan/decode.html?eJysk0Fv1DAQhe_8CmtOW2FpY29PlpCyXFChbFBVkBCKKhMPwSKJo7FDW1X731Hsw5Jss7CIW_w8n_3ejPMEnTO40y16UF9AQMmhJ1eh945GKRVcmQdQGQfb9UMY5ZJD5QhBPUGwoUFQcKu_NniD2iCtM-BgMGjbxGN7sq2mx_zHT-BA7t4zQm0UE8DBB900LNgWFcs8cCiGoFguoNxzcEM4XOiDrhGU2PN_MyUWTN3_B1Ny0dTBy9A5MkhoJj7KkfxTyTPJ3iPV-NbZDmktp8ka_BZWuXh58Yps_T19Ao8ymwWN2lHaiM1Lk3hU2-oH1mLr6JENHo1iMmPv7OtDz3guF9u2OWeW27omrHVwtN5MA-ejvyJ1Lo6Jw3b3-W5X3N7tPl5fr3Ix5t9-erPK5cVfTPsok0iZlmJcnhPjBn3vOo_zV_Dsydk4ejQ1pqfk3UAVfiBXxWvSsohcFAz6kHZlWlx1cSv-M7_D4gxYzmF5Et5M4GwOb07ClzO43L_4FQAA__-EfYCR
 
 # This query verifies stats collection for the hashJoiner, distinct and sorter.
 query T