executor: Support parallel sort spill

pkg/executor/sortexec/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "sortexec",
     srcs = [
+        "parallel_sort_spill_helper.go",
         "parallel_sort_worker.go",
         "sort.go",
         "sort_partition.go",
@@ -36,7 +37,7 @@ go_test(
     timeout = "short",
     srcs = ["sort_test.go"],
     flaky = True,
-    shard_count = 8,
+    shard_count = 10,
     deps = [
         "//pkg/config",
         "//pkg/sessionctx/variable",
@@ -52,6 +53,7 @@ go_test(
     timeout = "short",
     srcs = [
         "benchmark_test.go",
+        "parallel_sort_spill_test.go",
         "parallel_sort_test.go",
         "sort_spill_test.go",
         "sort_test.go",

pkg/executor/sortexec/parallel_sort_spill_helper.go

@@ -0,0 +1,225 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sortexec
+
+import (
+	"sync"
+	"sync/atomic"
+
+	"github.com/pingcap/tidb/pkg/types"
+	"github.com/pingcap/tidb/pkg/util"
+	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/logutil"
+	"go.uber.org/zap"
+)
+
+type parallelSortSpillHelper struct {
+	cond             *sync.Cond
+	spillStatus      int
+	sortedRowsInDisk []*chunk.DataInDiskByChunks
+	sortExec         *SortExec
+
+	lessRowFunc   func(chunk.Row, chunk.Row) int
+	errOutputChan chan rowWithError
+	finishCh      chan struct{}
+
+	fieldTypes    []*types.FieldType
+	tmpSpillChunk *chunk.Chunk
+
+	bytesConsumed atomic.Int64
+	bytesLimit    atomic.Int64
+}
+
+func newParallelSortSpillHelper(sortExec *SortExec, fieldTypes []*types.FieldType, finishCh chan struct{}, lessRowFunc func(chunk.Row, chunk.Row) int, errOutputChan chan rowWithError) *parallelSortSpillHelper {
+	return &parallelSortSpillHelper{
+		cond:          sync.NewCond(new(sync.Mutex)),
+		spillStatus:   notSpilled,
+		sortExec:      sortExec,
+		lessRowFunc:   lessRowFunc,
+		errOutputChan: errOutputChan,
+		finishCh:      finishCh,
+		fieldTypes:    fieldTypes,
+		tmpSpillChunk: chunk.NewChunkWithCapacity(fieldTypes, spillChunkSize),
+	}
+}
+
+func (p *parallelSortSpillHelper) isNotSpilledNoLock() bool {
+	return p.spillStatus == notSpilled
+}
+
+func (p *parallelSortSpillHelper) isInSpillingNoLock() bool {
+	return p.spillStatus == inSpilling
+}
+
+func (p *parallelSortSpillHelper) isSpillNeeded() bool {
+	p.cond.L.Lock()
+	defer p.cond.L.Unlock()
+	return p.spillStatus == needSpill
+}
+
+func (p *parallelSortSpillHelper) isSpillTriggered() bool {
+	p.cond.L.Lock()
+	defer p.cond.L.Unlock()
+	return len(p.sortedRowsInDisk) > 0
+}
+
+func (p *parallelSortSpillHelper) setInSpilling() {
+	p.cond.L.Lock()
+	defer p.cond.L.Unlock()
+	p.spillStatus = inSpilling
+}
+
+func (p *parallelSortSpillHelper) setNeedSpillNoLock() {
+	p.spillStatus = needSpill
+}
+
+func (p *parallelSortSpillHelper) setNotSpilled() {
+	p.cond.L.Lock()
+	defer p.cond.L.Unlock()
+	p.spillStatus = notSpilled
+}
+
+func (p *parallelSortSpillHelper) spill() (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = util.GetRecoverError(r)
+		}
+	}()
+
+	select {
+	case <-p.finishCh:
+		return nil
+	default:
+	}
+
+	workerNum := len(p.sortExec.Parallel.workers)
+	workerWaiter := &sync.WaitGroup{}
+	workerWaiter.Add(workerNum)
+	sortedRowsIters := make([]*chunk.Iterator4Slice, workerNum)
+	for i := 0; i < workerNum; i++ {
+		go func(idx int) {
+			defer func() {
+				if r := recover(); r != nil {
+					processPanicAndLog(p.errOutputChan, r)
+				}
+				workerWaiter.Done()
+			}()
+
+			sortedRowsIters[idx] = chunk.NewIterator4Slice(nil)
+			sortedRowsIters[idx].Reset(p.sortExec.Parallel.workers[idx].sortLocalRows())
+			injectParallelSortRandomFail(200)
+		}(i)
+	}
+
+	workerWaiter.Wait()
+	p.setInSpilling()
+
+	// Spill is done, broadcast to wake up all sleep goroutines
+	defer p.cond.Broadcast()
+	defer p.setNotSpilled()
+
+	totalRows := 0
+	for i := range sortedRowsIters {
+		totalRows += sortedRowsIters[i].Len()
+	}
+
+	merger := newMultiWayMerger(sortedRowsIters, p.lessRowFunc)
+	merger.init()
+	return p.spillImpl(merger)
+}
+
+func (p *parallelSortSpillHelper) releaseMemory() {
+	totalReleasedMemory := int64(0)
+	for _, worker := range p.sortExec.Parallel.workers {
+		totalReleasedMemory += worker.totalMemoryUsage
+		worker.totalMemoryUsage = 0
+	}
+	p.sortExec.memTracker.Consume(-totalReleasedMemory)
+}
+
+func (p *parallelSortSpillHelper) spillTmpSpillChunk(inDisk *chunk.DataInDiskByChunks) error {
+	err := inDisk.Add(p.tmpSpillChunk)
+	if err != nil {
+		return err
+	}
+	p.tmpSpillChunk.Reset()
+	return nil
+}
+
+func (p *parallelSortSpillHelper) spillImpl(merger *multiWayMerger) error {
+	logutil.BgLogger().Info(spillInfo, zap.Int64("consumed", p.bytesConsumed.Load()), zap.Int64("quota", p.bytesLimit.Load()))
+	p.tmpSpillChunk.Reset()
+	inDisk := chunk.NewDataInDiskByChunks(p.fieldTypes)
+	inDisk.GetDiskTracker().AttachTo(p.sortExec.diskTracker)
+
+	spilledRowChannel := make(chan chunk.Row, 10000)
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				processPanicAndLog(p.errOutputChan, r)
+			}
+			close(spilledRowChannel)
+		}()
+
+		injectParallelSortRandomFail(200)
+
+		for {
+			row := merger.next()
+			if row.IsEmpty() {
+				break
+			}
+			spilledRowChannel <- row
+		}
+	}()
+
+	var (
+		row chunk.Row
+		ok  bool
+	)
+	for {
+		select {
+		case <-p.finishCh:
+			// We must wait the finish of the above goroutine,
+			// or p.errOutputChan may be closed in advandce.
+			<-spilledRowChannel
+			return nil
+		case row, ok = <-spilledRowChannel:
+			if !ok {
+				if p.tmpSpillChunk.NumRows() > 0 {
+					err := p.spillTmpSpillChunk(inDisk)
+					if err != nil {
+						return err
+					}
+				}
+
+				injectParallelSortRandomFail(200)
+
+				if inDisk.NumRows() > 0 {
+					p.sortedRowsInDisk = append(p.sortedRowsInDisk, inDisk)
+					p.releaseMemory()
+				}
+				return nil
+			}
+		}
+
+		p.tmpSpillChunk.AppendRow(row)
+		if p.tmpSpillChunk.IsFull() {
+			err := p.spillTmpSpillChunk(inDisk)
+			if err != nil {
+				return err
+			}
+		}
+	}
+}

pkg/executor/sortexec/parallel_sort_spill_test.go

@@ -0,0 +1,141 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sortexec_test
+
+import (
+	"testing"
+
+	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/pkg/executor/internal/testutil"
+	"github.com/pingcap/tidb/pkg/executor/sortexec"
+	"github.com/pingcap/tidb/pkg/expression"
+	"github.com/pingcap/tidb/pkg/util/memory"
+	"github.com/pingcap/tidb/pkg/util/mock"
+	"github.com/stretchr/testify/require"
+)
+
+var hardLimit1 = int64(100000)
+var hardLimit2 = hardLimit1 * 10
+
+func oneSpillCase(t *testing.T, ctx *mock.Context, exe *sortexec.SortExec, sortCase *testutil.SortCase, schema *expression.Schema, dataSource *testutil.MockDataSource) {
+	if exe == nil {
+		exe = buildSortExec(ctx, sortCase, dataSource)
+	}
+	dataSource.PrepareChunks()
+	resultChunks := executeSortExecutor(t, exe)
+
+	require.True(t, exe.IsSpillTriggeredInParallelSortForTest())
+	require.Equal(t, int64(sortCase.Rows), exe.GetSpilledRowNumInParallelSortForTest())
+
+	err := exe.Close()
+	require.NoError(t, err)
+
+	require.True(t, checkCorrectness(schema, exe, dataSource, resultChunks))
+}
+
+func inMemoryThenSpill(t *testing.T, ctx *mock.Context, exe *sortexec.SortExec, sortCase *testutil.SortCase, schema *expression.Schema, dataSource *testutil.MockDataSource) {
+	if exe == nil {
+		exe = buildSortExec(ctx, sortCase, dataSource)
+	}
+	dataSource.PrepareChunks()
+	resultChunks := executeSortExecutorAndManullyTriggerSpill(t, exe, hardLimit2, ctx.GetSessionVars().StmtCtx.MemTracker)
+
+	require.True(t, exe.IsSpillTriggeredInParallelSortForTest())
+	require.Greater(t, int64(sortCase.Rows), exe.GetSpilledRowNumInParallelSortForTest())
+	err := exe.Close()
+	require.NoError(t, err)
+
+	require.True(t, checkCorrectness(schema, exe, dataSource, resultChunks))
+}
+
+func failpointNoMemoryDataTest(t *testing.T, ctx *mock.Context, exe *sortexec.SortExec, sortCase *testutil.SortCase, schema *expression.Schema, dataSource *testutil.MockDataSource) {
+	if exe == nil {
+		exe = buildSortExec(ctx, sortCase, dataSource)
+	}
+	dataSource.PrepareChunks()
+	executeInFailpoint(t, exe, 0, nil)
+}
+
+func failpointDataInMemoryThenSpillTest(t *testing.T, ctx *mock.Context, exe *sortexec.SortExec, sortCase *testutil.SortCase, schema *expression.Schema, dataSource *testutil.MockDataSource) {
+	if exe == nil {
+		exe = buildSortExec(ctx, sortCase, dataSource)
+	}
+	dataSource.PrepareChunks()
+	executeInFailpoint(t, exe, hardLimit2, ctx.GetSessionVars().MemTracker)
+}
+
+func TestParallelSortSpillDisk(t *testing.T) {
+	sortexec.SetSmallSpillChunkSizeForTest()
+	ctx := mock.NewContext()
+	sortCase := &testutil.SortCase{Rows: 10000, OrderByIdx: []int{0, 1}, Ndvs: []int{0, 0}, Ctx: ctx}
+
+	failpoint.Enable("github.com/pingcap/tidb/pkg/executor/sortexec/SlowSomeWorkers", `return(true)`)
+	failpoint.Enable("github.com/pingcap/tidb/pkg/executor/sortexec/SignalCheckpointForSort", `return(true)`)
+
+	ctx.GetSessionVars().InitChunkSize = 32
+	ctx.GetSessionVars().MaxChunkSize = 32
+	ctx.GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSQLText, hardLimit1)
+	ctx.GetSessionVars().StmtCtx.MemTracker = memory.NewTracker(memory.LabelForSQLText, -1)
+	ctx.GetSessionVars().StmtCtx.MemTracker.AttachTo(ctx.GetSessionVars().MemTracker)
+	ctx.GetSessionVars().EnableParallelSort = true
+
+	schema := expression.NewSchema(sortCase.Columns()...)
+	dataSource := buildDataSource(ctx, sortCase, schema)
+	exe := buildSortExec(ctx, sortCase, dataSource)
+	for i := 0; i < 10; i++ {
+		oneSpillCase(t, ctx, nil, sortCase, schema, dataSource)
+		oneSpillCase(t, ctx, exe, sortCase, schema, dataSource)
+	}
+
+	ctx.GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSQLText, hardLimit2)
+	ctx.GetSessionVars().StmtCtx.MemTracker.AttachTo(ctx.GetSessionVars().MemTracker)
+	for i := 0; i < 10; i++ {
+		inMemoryThenSpill(t, ctx, nil, sortCase, schema, dataSource)
+		inMemoryThenSpill(t, ctx, exe, sortCase, schema, dataSource)
+	}
+}
+
+func TestParallelSortSpillDiskFailpoint(t *testing.T) {
+	sortexec.SetSmallSpillChunkSizeForTest()
+	ctx := mock.NewContext()
+	sortCase := &testutil.SortCase{Rows: 10000, OrderByIdx: []int{0, 1}, Ndvs: []int{0, 0}, Ctx: ctx}
+
+	failpoint.Enable("github.com/pingcap/tidb/pkg/executor/sortexec/SlowSomeWorkers", `return(true)`)
+	failpoint.Enable("github.com/pingcap/tidb/pkg/executor/sortexec/SignalCheckpointForSort", `return(true)`)
+	failpoint.Enable("github.com/pingcap/tidb/pkg/executor/sortexec/ParallelSortRandomFail", `return(true)`)
+	failpoint.Enable("github.com/pingcap/tidb/pkg/util/chunk/ChunkInDiskError", `return(false)`)
+
+	ctx.GetSessionVars().InitChunkSize = 32
+	ctx.GetSessionVars().MaxChunkSize = 32
+	ctx.GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSQLText, hardLimit1)
+	ctx.GetSessionVars().StmtCtx.MemTracker = memory.NewTracker(memory.LabelForSQLText, -1)
+	ctx.GetSessionVars().StmtCtx.MemTracker.AttachTo(ctx.GetSessionVars().MemTracker)
+	ctx.GetSessionVars().EnableParallelSort = true
+
+	schema := expression.NewSchema(sortCase.Columns()...)
+	dataSource := buildDataSource(ctx, sortCase, schema)
+	exe := buildSortExec(ctx, sortCase, dataSource)
+	for i := 0; i < 20; i++ {
+		failpointNoMemoryDataTest(t, ctx, nil, sortCase, schema, dataSource)
+		failpointNoMemoryDataTest(t, ctx, exe, sortCase, schema, dataSource)
+	}
+
+	ctx.GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSQLText, hardLimit2)
+	ctx.GetSessionVars().StmtCtx.MemTracker.AttachTo(ctx.GetSessionVars().MemTracker)
+	for i := 0; i < 20; i++ {
+		failpointDataInMemoryThenSpillTest(t, ctx, nil, sortCase, schema, dataSource)
+		failpointDataInMemoryThenSpillTest(t, ctx, exe, sortCase, schema, dataSource)
+	}
+}