*: support canceling DDL statements with KILL (#35803)

close #24144

ddl/ddl.go

@@ -737,6 +737,11 @@ func (d *ddl) DoDDLJob(ctx sessionctx.Context, job *model.Job) error {
 
 	var historyJob *model.Job
 	jobID := job.ID
+
+	// Attach the context of the jobId to the calling session so that
+	// KILL can cancel this DDL job.
+	ctx.GetSessionVars().StmtCtx.DDLJobID = jobID
+
 	// For a job from start to end, the state of it will be none -> delete only -> write only -> reorganization -> public
 	// For every state changes, we will wait as lease 2 * lease time, so here the ticker check is 10 * lease.
 	// But we use etcd to speed up, normally it takes less than 0.5s now, so we use 0.5s or 1s or 3s as the max value.

executor/ddl.go

@@ -132,7 +132,10 @@ func (e *DDLExec) Next(ctx context.Context, req *chunk.Chunk) (err error) {
 		return err
 	}
 
-	defer func() { e.ctx.GetSessionVars().StmtCtx.IsDDLJobInQueue = false }()
+	defer func() {
+		e.ctx.GetSessionVars().StmtCtx.IsDDLJobInQueue = false
+		e.ctx.GetSessionVars().StmtCtx.DDLJobID = 0
+	}()
 
 	switch x := e.stmt.(type) {
 	case *ast.AlterDatabaseStmt:

server/server.go

@@ -50,6 +50,7 @@ import (
 	"github.com/blacktear23/go-proxyprotocol"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/ddl"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/kv"
@@ -718,6 +719,25 @@ func killConn(conn *clientConn) {
 	conn.mu.RLock()
 	cancelFunc := conn.mu.cancelFunc
 	conn.mu.RUnlock()
+
+	// If the connection being killed is a DDL Job,
+	// we need to CANCEL the matching jobID first.
+	if sessVars.StmtCtx.IsDDLJobInQueue {
+		jobID := sessVars.StmtCtx.DDLJobID
+		err := kv.RunInNewTxn(context.Background(), conn.ctx.GetStore(), true, func(ctx context.Context, txn kv.Transaction) error {
+			// errs is the error per job, there is only one submitted
+			// err is the error of the overall task
+			errs, err := ddl.CancelJobs(txn, []int64{jobID})
+			if len(errs) > 0 {
+				logutil.BgLogger().Warn("error canceling DDL job", zap.Error(errs[0]))
+			}
+			return err
+		})
+		if err != nil {
+			logutil.BgLogger().Warn("could not cancel DDL job", zap.Error(err))
+		}
+	}
+
 	if cancelFunc != nil {
 		cancelFunc()
 	}

sessionctx/stmtctx/stmtctx.go

@@ -109,6 +109,7 @@ type StatementContext struct {
 	// IsDDLJobInQueue is used to mark whether the DDL job is put into the queue.
 	// If IsDDLJobInQueue is true, it means the DDL job is in the queue of storage, and it can be handled by the DDL worker.
 	IsDDLJobInQueue        bool
+	DDLJobID               int64
 	InInsertStmt           bool
 	InUpdateStmt           bool
 	InDeleteStmt           bool