executor: fix goroutine leak when exceed quota in hash agg (#58078) (#61803)

close #58004

Author: ti-chi-bot

pkg/executor/aggregate/agg_hash_executor.go

@@ -154,6 +154,8 @@ type HashAggExec struct {
 	spillHelper *parallelHashAggSpillHelper
 	// isChildDrained indicates whether the all data from child has been taken out.
 	isChildDrained bool
+
+	invalidMemoryUsageForTrackingTest bool
 }
 
 // Close implements the Executor Close interface.
@@ -204,7 +206,11 @@ func (e *HashAggExec) Close() error {
 		channel.Clear(e.finalOutputCh)
 		e.executed.Store(false)
 		if e.memTracker != nil {
-			e.memTracker.ReplaceBytesUsed(0)
+			if e.memTracker.BytesConsumed() < 0 {
+				e.invalidMemoryUsageForTrackingTest = true
+			} else {
+				e.memTracker.ReplaceBytesUsed(0)
+			}
 		}
 		e.parallelExecValid = false
 		if e.parallelAggSpillAction != nil {
@@ -285,6 +291,8 @@ func (e *HashAggExec) initForUnparallelExec() {
 }
 
 func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrency int, ctx sessionctx.Context) {
+	memUsage := int64(0)
+
 	for i := 0; i < partialConcurrency; i++ {
 		partialResultsMap := make([]aggfuncs.AggPartialResultMapper, finalConcurrency)
 		for i := 0; i < finalConcurrency; i++ {
@@ -311,6 +319,8 @@ func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrenc
 			inflightChunkSync:    e.inflightChunkSync,
 		}
 
+		memUsage += e.partialWorkers[i].chk.MemoryUsage()
+
 		e.partialWorkers[i].partialResultNumInRow = e.partialWorkers[i].getPartialResultSliceLenConsiderByteAlign()
 		for j := 0; j < finalConcurrency; j++ {
 			e.partialWorkers[i].BInMaps[j] = 0
@@ -328,9 +338,11 @@ func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrenc
 			chk:        exec.NewFirstChunk(e.Children(0)),
 			giveBackCh: e.partialWorkers[i].inputCh,
 		}
-		e.memTracker.Consume(input.chk.MemoryUsage())
+		memUsage += input.chk.MemoryUsage()
 		e.inputCh <- input
 	}
+
+	e.memTracker.Consume(memUsage)
 }
 
 func (e *HashAggExec) initFinalWorkers(finalConcurrency int) {
@@ -442,6 +454,7 @@ func (e *HashAggExec) fetchChildData(ctx context.Context, waitGroup *sync.WaitGr
 		ok    bool
 		err   error
 	)
+
 	defer func() {
 		if r := recover(); r != nil {
 			recoveryHashAgg(e.finalOutputCh, r)
@@ -494,6 +507,7 @@ func (e *HashAggExec) fetchChildData(ctx context.Context, waitGroup *sync.WaitGr
 		input.giveBackCh <- chk
 
 		if hasError := e.spillIfNeed(); hasError {
+			e.memTracker.Consume(-mSize)
 			return
 		}
 	}
@@ -857,3 +871,8 @@ func (e *HashAggExec) IsSpillTriggeredForTest() bool {
 	}
 	return false
 }
+
+// IsInvalidMemoryUsageTrackingForTest is for test
+func (e *HashAggExec) IsInvalidMemoryUsageTrackingForTest() bool {
+	return e.invalidMemoryUsageForTrackingTest
+}

pkg/executor/aggregate/agg_spill_test.go

@@ -150,6 +150,7 @@ func generateResult(t *testing.T, ctx *mock.Context, dataSource *testutil.MockDa
 			resultRows = append(resultRows, chk.GetRow(i))
 		}
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 
 	require.False(t, aggExec.IsSpillTriggeredForTest())
@@ -315,6 +316,7 @@ func executeCorrecResultTest(t *testing.T, ctx *mock.Context, aggExec *aggregate
 			resultRows = append(resultRows, chk.GetRow(i))
 		}
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 
 	require.True(t, aggExec.IsSpillTriggeredForTest())
@@ -351,6 +353,7 @@ func fallBackActionTest(t *testing.T) {
 		}
 		chk.Reset()
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 	require.Less(t, 0, newRootExceedAction.GetTriggeredNum())
 }
@@ -373,6 +376,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 	go func() {
 		time.Sleep(time.Duration(rand.Int31n(300)) * time.Millisecond)
 		once.Do(func() {
+			require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 			aggExec.Close()
 		})
 		goRoutineWaiter.Done()
@@ -382,6 +386,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 		err := aggExec.Next(tmpCtx, chk)
 		if err != nil {
 			once.Do(func() {
+				require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 				err = aggExec.Close()
 				require.Equal(t, nil, err)
 			})
@@ -393,6 +398,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 		chk.Reset()
 	}
 	once.Do(func() {
+		require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 		aggExec.Close()
 	})
 }