From 97bb48b52ba4b777e3b85316b8bbbf867f9ca2d7 Mon Sep 17 00:00:00 2001
From: Feng Liyuan <darktemplar.f@gmail.com>
Date: Fri, 1 Nov 2019 16:18:04 +0800
Subject: [PATCH 1/2] server: fix potential goroutine leak when kill connection
 (#13051)

---
 server/conn.go      | 21 ++++++++++++---------
 server/conn_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/server/conn.go b/server/conn.go
index e74dbd42a1546..92394806bac9f 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -39,7 +39,6 @@ import (
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
-	"github.com/pingcap/tidb/plugin"
 	"io"
 	"net"
 	"runtime"
@@ -48,6 +47,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/pingcap/tidb/plugin"
+
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
@@ -947,21 +948,23 @@ func (cc *clientConn) handleLoadStats(ctx context.Context, loadStatsInfo *execut
 // There is a special query `load data` that does not return result, which is handled differently.
 // Query `load stats` does not return result either.
 func (cc *clientConn) handleQuery(ctx context.Context, sql string) (err error) {
-	rs, err := cc.ctx.Execute(ctx, sql)
+	rss, err := cc.ctx.Execute(ctx, sql)
 	if err != nil {
 		metrics.ExecuteErrorCounter.WithLabelValues(metrics.ExecuteErrorToLabel(err)).Inc()
 		return errors.Trace(err)
 	}
 	status := atomic.LoadInt32(&cc.status)
-	if status == connStatusShutdown || status == connStatusWaitShutdown {
-		killConn(cc)
-		return errors.New("killed by another connection")
+	if rss != nil && (status == connStatusShutdown || status == connStatusWaitShutdown) {
+		for _, rs := range rss {
+			terror.Call(rs.Close)
+		}
+		return executor.ErrQueryInterrupted
 	}
-	if rs != nil {
-		if len(rs) == 1 {
-			err = cc.writeResultset(ctx, rs[0], false, 0, 0)
+	if rss != nil {
+		if len(rss) == 1 {
+			err = cc.writeResultset(ctx, rss[0], false, 0, 0)
 		} else {
-			err = cc.writeMultiResultset(ctx, rs, false)
+			err = cc.writeMultiResultset(ctx, rss, false)
 		}
 	} else {
 		loadDataInfo := cc.ctx.Value(executor.LoadDataVarKey)
diff --git a/server/conn_test.go b/server/conn_test.go
index 4ada64fb5568b..f439e7324b7f0 100644
--- a/server/conn_test.go
+++ b/server/conn_test.go
@@ -22,10 +22,14 @@ import (
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/mysql"
+	"github.com/pingcap/tidb/executor"
 	"github.com/pingcap/tidb/session"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/store/mockstore"
 	"github.com/pingcap/tidb/util/arena"
+	"github.com/pingcap/tidb/util/chunk"
 )
 
 type ConnTestSuite struct{}
@@ -235,3 +239,38 @@ func (ts ConnTestSuite) TestConnExecutionTimeout(c *C) {
 
 	c.Assert(failpoint.Disable("github.com/pingcap/tidb/server/FakeClientConn"), IsNil)
 }
+
+type mockTiDBCtx struct {
+	TiDBContext
+	rs  []ResultSet
+	err error
+}
+
+func (c *mockTiDBCtx) Execute(ctx context.Context, sql string) ([]ResultSet, error) {
+	return c.rs, c.err
+}
+
+func (c *mockTiDBCtx) GetSessionVars() *variable.SessionVars {
+	return &variable.SessionVars{}
+}
+
+type mockRecordSet struct{}
+
+func (m mockRecordSet) Fields() []*ast.ResultField                       { return nil }
+func (m mockRecordSet) Next(ctx context.Context, req *chunk.Chunk) error { return nil }
+func (m mockRecordSet) NewChunk() *chunk.Chunk                           { return nil }
+func (m mockRecordSet) Close() error                                     { return nil }
+
+func (ts *ConnTestSuite) TestShutDown(c *C) {
+	cc := &clientConn{}
+
+	rs := &tidbResultSet{recordSet: mockRecordSet{}}
+	// mock delay response
+	cc.ctx = &mockTiDBCtx{rs: []ResultSet{rs}, err: nil}
+	// set killed flag
+	cc.status = connStatusShutdown
+	// assert ErrQueryInterrupted
+	err := cc.handleQuery(context.Background(), "dummy")
+	c.Assert(err, Equals, executor.ErrQueryInterrupted)
+	c.Assert(rs.closed, Equals, int32(1))
+}

From b5566335a42272eb246fbdea9d1f7cd00e1663b0 Mon Sep 17 00:00:00 2001
From: Feng Liyuan <darktemplar.f@gmail.com>
Date: Thu, 7 Nov 2019 19:44:31 +0800
Subject: [PATCH 2/2] fix import

---
 server/conn.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/server/conn.go b/server/conn.go
index 92394806bac9f..ce99a05a8d441 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -47,8 +47,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/pingcap/tidb/plugin"
-
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
@@ -58,6 +56,7 @@ import (
 	"github.com/pingcap/tidb/executor"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/metrics"
+	"github.com/pingcap/tidb/plugin"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/arena"