store/tikv: call Next() after copIterator closed lead to goroutine leak (#5624)

executor/executor_test.go

@@ -17,6 +17,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -2306,3 +2307,36 @@ func (s *testSuite) TestTableScanWithPointRanges(c *C) {
 	tk.MustExec("insert into t values(1), (5), (10)")
 	tk.MustQuery("select * from t where id in(1, 2, 10)").Check(testkit.Rows("1", "10"))
 }
+
+func (s *testSuite) TestEarlyClose(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("create table earlyclose (id int primary key)")
+
+	// Insert 1000 rows.
+	var values []string
+	for i := 0; i < 1000; i++ {
+		values = append(values, fmt.Sprintf("(%d)", i))
+	}
+	tk.MustExec("insert earlyclose values " + strings.Join(values, ","))
+
+	// Get table ID for split.
+	dom := domain.GetDomain(tk.Se)
+	is := dom.InfoSchema()
+	tbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("earlyclose"))
+	c.Assert(err, IsNil)
+	tblID := tbl.Meta().ID
+
+	// Split the table.
+	s.cluster.SplitTable(s.mvccStore, tblID, 500)
+
+	goCtx := goctx.Background()
+	for i := 0; i < 500; i++ {
+		rss, err := tk.Se.Execute(goCtx, "select * from earlyclose order by id")
+		c.Assert(err, IsNil)
+		rs := rss[0]
+		_, err = rs.Next(goCtx)
+		c.Assert(err, IsNil)
+		rs.Close()
+	}
+}

store/tikv/coprocessor.go

@@ -434,6 +434,15 @@ func (it *copIterator) run(ctx goctx.Context) {
 	})
 }
 
+func recvFromRespCh(respCh <-chan copResponse, finished <-chan struct{}) (resp copResponse, ok bool, exit bool) {
+	select {
+	case resp, ok = <-respCh:
+	case <-finished:
+		exit = true
+	}
+	return
+}
+
 func (it *copIterator) sendToTaskCh(ctx goctx.Context, t *copTask, taskCh chan<- *copTask) (exit bool) {
 	select {
 	case taskCh <- t:
@@ -475,13 +484,18 @@ func (it *copIterator) Next() ([]byte, error) {
 			return nil, nil
 		}
 	} else {
+		var closed bool
 		for {
 			if it.curr >= len(it.tasks) {
 				// Resp will be nil if iterator is finished.
 				return nil, nil
 			}
 			task := it.tasks[it.curr]
-			resp, ok = <-task.respChan
+			resp, ok, closed = recvFromRespCh(task.respChan, it.finished)
+			if closed {
+				// Close() is already called, so Next() is invalid.
+				return nil, nil
+			}
 			if ok {
 				break
 			}