From 984487bcbba0bf1034a48cb9a7f6777dcf59249a Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 18 Aug 2022 14:17:03 -0700
Subject: [PATCH] colexec: fix a couple of places in sort where we
 under-account

Release justification: bug fix.

Release note: None
---
 pkg/sql/colexec/sort.eg.go   | 24 ++++++++++++++++++++----
 pkg/sql/colexec/sort.go      |  4 ++++
 pkg/sql/colexec/sort_tmpl.go |  6 +++++-
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/pkg/sql/colexec/sort.eg.go b/pkg/sql/colexec/sort.eg.go
index 5f9206c648d4..2623f7fd2a5f 100644
--- a/pkg/sql/colexec/sort.eg.go
+++ b/pkg/sql/colexec/sort.eg.go
@@ -399,13 +399,17 @@ func (s *sortBytesAscWithNullsOp) init(
 	s.allocator = allocator
 	s.allocator.AdjustMemoryUsage(memsize.Uint64 * int64(s.sortCol.Len()))
 	s.abbreviatedSortCol = s.sortCol.Abbreviated()
+	// Finalize the accounting in case the capacity of the new slice is larger
+	// than we asked for.
+	extraCap := cap(s.abbreviatedSortCol) - len(s.abbreviatedSortCol)
+	s.allocator.AdjustMemoryUsageAfterAllocation(memsize.Uint64 * int64(extraCap))
 	s.nulls = col.Nulls()
 	s.order = order
 	s.cancelChecker.Init(ctx)
 }
 
 func (s *sortBytesAscWithNullsOp) reset() {
-	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(s.sortCol.Len()))
+	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(cap(s.abbreviatedSortCol)))
 	s.allocator = nil
 	s.abbreviatedSortCol = nil
 	s.sortCol = nil
@@ -1371,13 +1375,17 @@ func (s *sortBytesDescWithNullsOp) init(
 	s.allocator = allocator
 	s.allocator.AdjustMemoryUsage(memsize.Uint64 * int64(s.sortCol.Len()))
 	s.abbreviatedSortCol = s.sortCol.Abbreviated()
+	// Finalize the accounting in case the capacity of the new slice is larger
+	// than we asked for.
+	extraCap := cap(s.abbreviatedSortCol) - len(s.abbreviatedSortCol)
+	s.allocator.AdjustMemoryUsageAfterAllocation(memsize.Uint64 * int64(extraCap))
 	s.nulls = col.Nulls()
 	s.order = order
 	s.cancelChecker.Init(ctx)
 }
 
 func (s *sortBytesDescWithNullsOp) reset() {
-	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(s.sortCol.Len()))
+	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(cap(s.abbreviatedSortCol)))
 	s.allocator = nil
 	s.abbreviatedSortCol = nil
 	s.sortCol = nil
@@ -2333,13 +2341,17 @@ func (s *sortBytesAscOp) init(
 	s.allocator = allocator
 	s.allocator.AdjustMemoryUsage(memsize.Uint64 * int64(s.sortCol.Len()))
 	s.abbreviatedSortCol = s.sortCol.Abbreviated()
+	// Finalize the accounting in case the capacity of the new slice is larger
+	// than we asked for.
+	extraCap := cap(s.abbreviatedSortCol) - len(s.abbreviatedSortCol)
+	s.allocator.AdjustMemoryUsageAfterAllocation(memsize.Uint64 * int64(extraCap))
 	s.nulls = col.Nulls()
 	s.order = order
 	s.cancelChecker.Init(ctx)
 }
 
 func (s *sortBytesAscOp) reset() {
-	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(s.sortCol.Len()))
+	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(cap(s.abbreviatedSortCol)))
 	s.allocator = nil
 	s.abbreviatedSortCol = nil
 	s.sortCol = nil
@@ -3198,13 +3210,17 @@ func (s *sortBytesDescOp) init(
 	s.allocator = allocator
 	s.allocator.AdjustMemoryUsage(memsize.Uint64 * int64(s.sortCol.Len()))
 	s.abbreviatedSortCol = s.sortCol.Abbreviated()
+	// Finalize the accounting in case the capacity of the new slice is larger
+	// than we asked for.
+	extraCap := cap(s.abbreviatedSortCol) - len(s.abbreviatedSortCol)
+	s.allocator.AdjustMemoryUsageAfterAllocation(memsize.Uint64 * int64(extraCap))
 	s.nulls = col.Nulls()
 	s.order = order
 	s.cancelChecker.Init(ctx)
 }
 
 func (s *sortBytesDescOp) reset() {
-	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(s.sortCol.Len()))
+	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(cap(s.abbreviatedSortCol)))
 	s.allocator = nil
 	s.abbreviatedSortCol = nil
 	s.sortCol = nil
diff --git a/pkg/sql/colexec/sort.go b/pkg/sql/colexec/sort.go
index 85b5c77e8d6e..cb7cea980c67 100644
--- a/pkg/sql/colexec/sort.go
+++ b/pkg/sql/colexec/sort.go
@@ -336,6 +336,10 @@ func (p *sortOp) sort() {
 		sizeAfter := memsize.Int * int64(spooledTuples)
 		p.allocator.AdjustMemoryUsage(sizeAfter - sizeBefore)
 		p.order = make([]int, spooledTuples)
+		// Finalize the accounting in case the capacity of the new slice is
+		// larger than we asked for.
+		extraCap := cap(p.order) - len(p.order)
+		p.allocator.AdjustMemoryUsageAfterAllocation(memsize.Int * int64(extraCap))
 	}
 	order := p.order[:spooledTuples]
 
diff --git a/pkg/sql/colexec/sort_tmpl.go b/pkg/sql/colexec/sort_tmpl.go
index 5faa02e7e508..ce1223871188 100644
--- a/pkg/sql/colexec/sort_tmpl.go
+++ b/pkg/sql/colexec/sort_tmpl.go
@@ -122,6 +122,10 @@ func (s *sort_TYPE_DIR_HANDLES_NULLSOp) init(
 	s.allocator = allocator
 	s.allocator.AdjustMemoryUsage(memsize.Uint64 * int64(s.sortCol.Len()))
 	s.abbreviatedSortCol = s.sortCol.Abbreviated()
+	// Finalize the accounting in case the capacity of the new slice is larger
+	// than we asked for.
+	extraCap := cap(s.abbreviatedSortCol) - len(s.abbreviatedSortCol)
+	s.allocator.AdjustMemoryUsageAfterAllocation(memsize.Uint64 * int64(extraCap))
 	// {{end}}
 	s.nulls = col.Nulls()
 	s.order = order
@@ -130,7 +134,7 @@ func (s *sort_TYPE_DIR_HANDLES_NULLSOp) init(
 
 func (s *sort_TYPE_DIR_HANDLES_NULLSOp) reset() {
 	// {{if .CanAbbreviate}}
-	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(s.sortCol.Len()))
+	s.allocator.AdjustMemoryUsage(0 - memsize.Uint64*int64(cap(s.abbreviatedSortCol)))
 	s.allocator = nil
 	s.abbreviatedSortCol = nil
 	// {{end}}