executor: avoid goroutine leak in index Lookup join (#19251) (#20791)

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
pingcap · Nov 12, 2020 · 6af0e02 · 6af0e02
1 parent b157b89
commit 6af0e02
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/cmd/benchdb/main.go b/cmd/benchdb/main.go
@@ -111,8 +111,17 @@ func newBenchDB() *benchDB {
 
 func (ut *benchDB) mustExec(sql string) {
 	rss, err := ut.session.Execute(context.Background(), sql)
+	defer func() {
+		for _, rs := range rss {
+			err = rs.Close()
+			if err != nil {
+				log.Fatal(err.Error())
+			}
+		}
+	}()
 	if err != nil {
 		log.Fatal(err.Error())
+		return
 	}
 	if len(rss) > 0 {
 		ctx := context.Background()

diff --git a/executor/index_lookup_join.go b/executor/index_lookup_join.go
@@ -601,10 +601,12 @@ func compareRow(sc *stmtctx.StatementContext, left, right []types.Datum) int {
 
 func (iw *innerWorker) fetchInnerResults(ctx context.Context, task *lookUpJoinTask, lookUpContent []*indexJoinLookUpContent) error {
 	innerExec, err := iw.readerBuilder.buildExecutorForIndexJoin(ctx, lookUpContent, iw.indexRanges, iw.keyOff2IdxOff, iw.nextColCompareFilters)
+	if innerExec != nil {
+		defer terror.Call(innerExec.Close)
+	}
 	if err != nil {
 		return err
 	}
-	defer terror.Call(innerExec.Close)
 	innerResult := chunk.NewList(retTypes(innerExec), iw.ctx.GetSessionVars().MaxChunkSize, iw.ctx.GetSessionVars().MaxChunkSize)
 	innerResult.GetMemTracker().SetLabel(innerResultLabel)
 	innerResult.GetMemTracker().AttachTo(task.memTracker)

diff --git a/session/session.go b/session/session.go
@@ -819,6 +819,14 @@ func (s *session) ExecRestrictedSQLWithSnapshot(sctx sessionctx.Context, sql str
 func execRestrictedSQL(ctx context.Context, se *session, sql string) ([]chunk.Row, []*ast.ResultField, error) {
 	startTime := time.Now()
 	recordSets, err := se.Execute(ctx, sql)
+	defer func() {
+		for _, rs := range recordSets {
+			closeErr := rs.Close()
+			if closeErr != nil && err == nil {
+				err = closeErr
+			}
+		}
+	}()
 	if err != nil {
 		return nil, nil, err
 	}
@@ -833,17 +841,14 @@ func execRestrictedSQL(ctx context.Context, se *session, sql string) ([]chunk.Ro
 		if err != nil {
 			return nil, nil, err
 		}
-		if err = rs.Close(); err != nil {
-			return nil, nil, err
-		}
 
 		if i == 0 {
 			rows = tmp
 			fields = rs.Fields()
 		}
 	}
 	metrics.QueryDurationHistogram.WithLabelValues(metrics.LblInternal).Observe(time.Since(startTime).Seconds())
-	return rows, fields, nil
+	return rows, fields, err
 }
 
 func createSessionFunc(store kv.Storage) pools.Factory {