MPP: Kill mpp queries

distsql/distsql.go

@@ -30,7 +30,7 @@ import (
 
 // DispatchMPPTasks dispathes all tasks and returns an iterator.
 func DispatchMPPTasks(ctx context.Context, sctx sessionctx.Context, tasks []*kv.MPPDispatchRequest, fieldTypes []*types.FieldType, planIDs []int, rootID int) (SelectResult, error) {
- resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, tasks)
+ resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, sctx.GetSessionVars().KVVars, tasks)
  if resp == nil {
  err := errors.New("client returns nil response")
  return nil, err

executor/mpp_gather.go

@@ -82,6 +82,7 @@ func (e *MPPGather) appendMPPDispatchReq(pf *plannercore.Fragment, tasks []*kv.M
  Timeout: 10,
  SchemaVar: e.is.SchemaMetaVersion(),
  StartTs: e.startTS,
+ State: kv.MppTaskReady,
  }
  e.mppReqs = append(e.mppReqs, req)
  }

executor/tiflash_test.go

@@ -15,18 +15,25 @@ package executor_test
 
 import (
  "fmt"
+ "sync"
+ "sync/atomic"
+ "time"
 
  . "github.com/pingcap/check"
+ "github.com/pingcap/errors"
  "github.com/pingcap/failpoint"
  "github.com/pingcap/kvproto/pkg/metapb"
  "github.com/pingcap/parser"
+ "github.com/pingcap/parser/terror"
  "github.com/pingcap/tidb/domain"
+ "github.com/pingcap/tidb/executor"
  "github.com/pingcap/tidb/kv"
  "github.com/pingcap/tidb/session"
  "github.com/pingcap/tidb/store/mockstore"
  "github.com/pingcap/tidb/store/mockstore/unistore"
  "github.com/pingcap/tidb/store/tikv/mockstore/cluster"
  "github.com/pingcap/tidb/util/testkit"
+ "github.com/pingcap/tidb/util/testleak"
 )
 
 type tiflashTestSuite struct {
@@ -36,6 +43,7 @@ type tiflashTestSuite struct {
 }
 
 func (s *tiflashTestSuite) SetUpSuite(c *C) {
+ testleak.BeforeTest()
  var err error
  s.store, err = mockstore.NewMockStore(
  mockstore.WithClusterInspector(func(c cluster.Cluster) {
@@ -65,6 +73,7 @@ func (s *tiflashTestSuite) SetUpSuite(c *C) {
 }
 
 func (s *tiflashTestSuite) TestReadPartitionTable(c *C) {
+ defer testleak.AfterTest(c)()
  tk := testkit.NewTestKit(c, s.store)
  tk.MustExec("use test")
  tk.MustExec("drop table if exists t")
@@ -92,6 +101,7 @@ func (s *tiflashTestSuite) TestReadPartitionTable(c *C) {
 }
 
 func (s *tiflashTestSuite) TestReadUnsigedPK(c *C) {
+ defer testleak.AfterTest(c)()
  tk := testkit.NewTestKit(c, s.store)
  tk.MustExec("use test")
  tk.MustExec("drop table if exists t")
@@ -127,6 +137,7 @@ func (s *tiflashTestSuite) TestReadUnsigedPK(c *C) {
 }
 
 func (s *tiflashTestSuite) TestMppExecution(c *C) {
+ defer testleak.AfterTest(c)()
  tk := testkit.NewTestKit(c, s.store)
  tk.MustExec("use test")
  tk.MustExec("drop table if exists t")
@@ -185,6 +196,7 @@ func (s *tiflashTestSuite) TestMppExecution(c *C) {
 }
 
 func (s *tiflashTestSuite) TestPartitionTable(c *C) {
+ defer testleak.AfterTest(c)()
  tk := testkit.NewTestKit(c, s.store)
  tk.MustExec("use test")
  tk.MustExec("drop table if exists t")
@@ -271,3 +283,75 @@ func (s *tiflashTestSuite) TestPartitionTable(c *C) {
  failpoint.Disable("github.com/pingcap/tidb/executor/checkTotalMPPTasks")
  failpoint.Disable("github.com/pingcap/tidb/executor/checkUseMPP")
 }
+
+func (s *tiflashTestSuite) TestCancelMppTasks(c *C) {
+ defer testleak.AfterTest(c)()
+ var hang = "github.com/pingcap/tidb/store/mockstore/unistore/mppRecvHang"
+ tk := testkit.NewTestKit(c, s.store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t")
+ tk.MustExec("create table t(a int not null primary key, b int not null)")
+ tk.MustExec("alter table t set tiflash replica 1")
+ tk.MustExec("insert into t values(1,0)")
+ tk.MustExec("insert into t values(2,0)")
+ tk.MustExec("insert into t values(3,0)")
+ tk.MustExec("insert into t values(4,0)")
+ tb := testGetTableByName(c, tk.Se, "test", "t")
+ err := domain.GetDomain(tk.Se).DDL().UpdateTableReplicaInfo(tk.Se, tb.Meta().ID, true)
+ c.Assert(err, IsNil)
+ tk.MustExec("set @@session.tidb_isolation_read_engines=\"tiflash\"")
+ tk.MustExec("set @@session.tidb_allow_mpp=ON")
+ atomic.StoreUint32(&tk.Se.GetSessionVars().Killed, 0)
+ c.Assert(failpoint.Enable(hang, `return(true)`), IsNil)
+ wg := &sync.WaitGroup{}
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ err := tk.QueryToErr("select count(*) from t as t1 , t where t1.a = t.a")
+ c.Assert(err, NotNil)
+ c.Assert(int(terror.ToSQLError(errors.Cause(err).(*terror.Error)).Code), Equals, int(executor.ErrQueryInterrupted.Code()))
+ }()
+ time.Sleep(1 * time.Second)
+ atomic.StoreUint32(&tk.Se.GetSessionVars().Killed, 1)
+ wg.Wait()
+ c.Assert(failpoint.Disable(hang), IsNil)
+}
+
+// all goroutines exit if one goroutine hangs but another return errors
+func (s *tiflashTestSuite) TestMppGoroutinesExitFromErrors(c *C) {
+ defer testleak.AfterTest(c)()
+ // mock non-root tasks return error
+ var mppNonRootTaskError = "github.com/pingcap/tidb/store/copr/mppNonRootTaskError"
+ // mock root tasks hang
+ var hang = "github.com/pingcap/tidb/store/mockstore/unistore/mppRecvHang"
+ tk := testkit.NewTestKit(c, s.store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t")
+ tk.MustExec("create table t(a int not null primary key, b int not null)")
+ tk.MustExec("alter table t set tiflash replica 1")
+ tb := testGetTableByName(c, tk.Se, "test", "t")
+ err := domain.GetDomain(tk.Se).DDL().UpdateTableReplicaInfo(tk.Se, tb.Meta().ID, true)
+ c.Assert(err, IsNil)
+ tk.MustExec("insert into t values(1,0)")
+ tk.MustExec("insert into t values(2,0)")
+ tk.MustExec("insert into t values(3,0)")
+ tk.MustExec("drop table if exists t1")
+ tk.MustExec("create table t1(a int not null primary key, b int not null)")
+ tk.MustExec("alter table t1 set tiflash replica 1")
+ tb = testGetTableByName(c, tk.Se, "test", "t1")
+ err = domain.GetDomain(tk.Se).DDL().UpdateTableReplicaInfo(tk.Se, tb.Meta().ID, true)
+ c.Assert(err, IsNil)
+ tk.MustExec("insert into t1 values(1,0)")
+ tk.MustExec("insert into t1 values(2,0)")
+ tk.MustExec("insert into t1 values(3,0)")
+ tk.MustExec("set @@session.tidb_isolation_read_engines=\"tiflash\"")
+ tk.MustExec("set @@session.tidb_allow_mpp=ON")
+ c.Assert(failpoint.Enable(mppNonRootTaskError, `return(true)`), IsNil)
+ c.Assert(failpoint.Enable(hang, `return(true)`), IsNil)
+
+ // generate 2 root tasks, one will hang and another will return errors
+ err = tk.QueryToErr("select count(*) from t as t1 , t where t1.a = t.a")
+ c.Assert(err, NotNil)
+ c.Assert(failpoint.Disable(mppNonRootTaskError), IsNil)
+ c.Assert(failpoint.Disable(hang), IsNil)
+}

go.mod

@@ -31,7 +31,7 @@ require (
  github.com/kr/text v0.2.0 // indirect
  github.com/ngaut/pools v0.0.0-20180318154953-b7bc8c42aac7
  github.com/ngaut/sync2 v0.0.0-20141008032647-7a24ed77b2ef
- github.com/ngaut/unistore v0.0.0-20210304095907-0ebafaf44efb
+ github.com/ngaut/unistore v0.0.0-20210310131351-7ad6a204de87
  github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
  github.com/opentracing/basictracer-go v1.0.0
  github.com/opentracing/opentracing-go v1.1.0
@@ -80,7 +80,7 @@ require (
  gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
  gopkg.in/natefinch/lumberjack.v2 v2.0.0
  gopkg.in/yaml.v2 v2.3.0 // indirect
- honnef.co/go/tools v0.1.2 // indirect
+ honnef.co/go/tools v0.1.3 // indirect
  sourcegraph.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0
  sourcegraph.com/sourcegraph/appdash-data v0.0.0-20151005221446-73f23eafcf67
 )

go.sum

@@ -353,8 +353,8 @@ github.com/ngaut/pools v0.0.0-20180318154953-b7bc8c42aac7 h1:7KAv7KMGTTqSmYZtNdc
 github.com/ngaut/pools v0.0.0-20180318154953-b7bc8c42aac7/go.mod h1:iWMfgwqYW+e8n5lC/jjNEhwcjbRDpl5NT7n2h+4UNcI=
 github.com/ngaut/sync2 v0.0.0-20141008032647-7a24ed77b2ef h1:K0Fn+DoFqNqktdZtdV3bPQ/0cuYh2H4rkg0tytX/07k=
 github.com/ngaut/sync2 v0.0.0-20141008032647-7a24ed77b2ef/go.mod h1:7WjlapSfwQyo6LNmIvEWzsW1hbBQfpUO4JWnuQRmva8=
-github.com/ngaut/unistore v0.0.0-20210304095907-0ebafaf44efb h1:2rGvEhflp/uK1l1rNUmoHA4CiHpbddHGxg52H71Fke8=
-github.com/ngaut/unistore v0.0.0-20210304095907-0ebafaf44efb/go.mod h1:ZR3NH+HzqfiYetwdoAivApnIy8iefPZHTMLfrFNm8g4=
+github.com/ngaut/unistore v0.0.0-20210310131351-7ad6a204de87 h1:lVRrhmqIT2zMbmoalrgxQLwWzFd3VtFaaWy0fnMwPro=
+github.com/ngaut/unistore v0.0.0-20210310131351-7ad6a204de87/go.mod h1:ZR3NH+HzqfiYetwdoAivApnIy8iefPZHTMLfrFNm8g4=
 github.com/nicksnyder/go-i18n v1.10.0/go.mod h1:HrK7VCrbOvQoUAQ7Vpy7i87N7JZZZ7R2xBGjv0j365Q=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
@@ -845,9 +845,10 @@ honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
+honnef.co/go/tools v0.0.1-2019.2.3 h1:3JgtbtFHMiCmsznwGVTUWbgGov+pVqnlf1dEJTNAXeM=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
-honnef.co/go/tools v0.1.2 h1:SMdYLJl312RXuxXziCCHhRsp/tvct9cGKey0yv95tZM=
-honnef.co/go/tools v0.1.2/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las=
+honnef.co/go/tools v0.1.3 h1:qTakTkI6ni6LFD5sBwwsdSO+AQqbSIxOauHTTQKZ/7o=
+honnef.co/go/tools v0.1.3/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las=
 k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 sigs.k8s.io/yaml v1.1.0 h1:4A07+ZFc2wgJwo8YNlQpr1rVlgUDlxXHhPJciaPY5gs=

kv/mpp.go

@@ -45,6 +45,20 @@ func (t *MPPTask) ToPB() *mpp.TaskMeta {
  return meta
 }
 
+//MppTaskStates denotes the state of mpp tasks
+type MppTaskStates uint8
+
+const (
+ // MppTaskReady means the task is ready
+ MppTaskReady MppTaskStates = iota
+ // MppTaskRunning means the task is running
+ MppTaskRunning
+ // MppTaskCancelled means the task is cancelled
+ MppTaskCancelled
+ // MppTaskDone means the task is done
+ MppTaskDone
+)
+
 // MPPDispatchRequest stands for a dispatching task.
 type MPPDispatchRequest struct {
  Data []byte // data encodes the dag coprocessor request.
@@ -55,6 +69,7 @@ type MPPDispatchRequest struct {
  SchemaVar int64
  StartTs uint64
  ID int64 // identify a single task
+ State MppTaskStates
 }
 
 // MPPClient accepts and processes mpp requests.
@@ -64,7 +79,7 @@ type MPPClient interface {
  ConstructMPPTasks(context.Context, *MPPBuildTasksRequest) ([]MPPTaskMeta, error)
 
  // DispatchMPPTasks dispatches ALL mpp requests at once, and returns an iterator that transfers the data.
- DispatchMPPTasks(context.Context, []*MPPDispatchRequest) Response
+ DispatchMPPTasks(context.Context, *Variables, []*MPPDispatchRequest) Response
 }
 
 // MPPBuildTasksRequest request the stores allocation for a mpp plan fragment.