From e4f5426a4d6b2bf54a96cb7ce89a7575877f623e Mon Sep 17 00:00:00 2001 From: pingcap-github-bot Date: Mon, 30 Mar 2020 19:17:42 +0800 Subject: [PATCH] distsql: check the killed flag in the distsql package (#15592) (#15616) --- distsql/select_result.go | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/distsql/select_result.go b/distsql/select_result.go index 843e8b048ba61..d03ee2838f6cb 100644 --- a/distsql/select_result.go +++ b/distsql/select_result.go @@ -16,9 +16,11 @@ package distsql import ( "context" "fmt" + "sync/atomic" "time" "github.com/pingcap/errors" + "github.com/pingcap/parser/mysql" "github.com/pingcap/parser/terror" "github.com/pingcap/tidb/kv" "github.com/pingcap/tidb/metrics" @@ -33,6 +35,10 @@ import ( "go.uber.org/zap" ) +var ( + errQueryInterrupted = terror.ClassExecutor.New(mysql.ErrQueryInterrupted, mysql.MySQLErrName[mysql.ErrQueryInterrupted]) +) + var ( _ SelectResult = (*selectResult)(nil) _ SelectResult = (*streamResult)(nil) @@ -188,7 +194,11 @@ func (r *selectResult) getSelectResp() error { if err := r.selectResp.Error; err != nil { return terror.ClassTiKV.New(terror.ErrCode(err.Code), err.Msg) } - sc := r.ctx.GetSessionVars().StmtCtx + sessVars := r.ctx.GetSessionVars() + if atomic.CompareAndSwapUint32(&sessVars.Killed, 1, 0) { + return errors.Trace(errQueryInterrupted) + } + sc := sessVars.StmtCtx for _, warning := range r.selectResp.Warnings { sc.AppendWarning(terror.ClassTiKV.New(terror.ErrCode(warning.Code), warning.Msg)) }