executor,server: re-implement the kill statement by checking the Next() function

executor/adapter.go

@@ -106,7 +106,7 @@ func (a *recordSet) Next(ctx context.Context, req *chunk.RecordBatch) error {
 		defer span1.Finish()
 	}
 
-	err := a.executor.Next(ctx, req)
+	err := Next(ctx, a.executor, req)
 	if err != nil {
 		a.lastErr = err
 		return err
@@ -385,7 +385,7 @@ func (a *ExecStmt) handleNoDelayExecutor(ctx context.Context, e Executor) (sqlex
 		a.logAudit()
 	}()
 
-	err = e.Next(ctx, chunk.NewRecordBatch(e.newFirstChunk()))
+	err = Next(ctx, e, chunk.NewRecordBatch(e.newFirstChunk()))
 	if err != nil {
 		return nil, err
 	}

executor/aggregate.go

@@ -555,7 +555,7 @@ func (e *HashAggExec) fetchChildData(ctx context.Context) {
 			}
 			chk = input.chk
 		}
-		err = e.children[0].Next(ctx, chunk.NewRecordBatch(chk))
+		err = Next(ctx, e.children[0], chunk.NewRecordBatch(chk))
 		if err != nil {
 			e.finalOutputCh <- &AfFinalResult{err: err}
 			return
@@ -681,7 +681,7 @@ func (e *HashAggExec) unparallelExec(ctx context.Context, chk *chunk.Chunk) erro
 func (e *HashAggExec) execute(ctx context.Context) (err error) {
 	inputIter := chunk.NewIterator4Chunk(e.childResult)
 	for {
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(e.childResult))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(e.childResult))
 		if err != nil {
 			return err
 		}
@@ -870,7 +870,7 @@ func (e *StreamAggExec) fetchChildIfNecessary(ctx context.Context, chk *chunk.Ch
 		return err
 	}
 
-	err = e.children[0].Next(ctx, chunk.NewRecordBatch(e.childResult))
+	err = Next(ctx, e.children[0], chunk.NewRecordBatch(e.childResult))
 	if err != nil {
 		return err
 	}

executor/delete.go

@@ -105,7 +105,7 @@ func (e *DeleteExec) deleteSingleTableByChunk(ctx context.Context) error {
 	for {
 		iter := chunk.NewIterator4Chunk(chk)
 
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(chk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(chk))
 		if err != nil {
 			return err
 		}
@@ -187,7 +187,7 @@ func (e *DeleteExec) deleteMultiTablesByChunk(ctx context.Context) error {
 	chk := e.children[0].newFirstChunk()
 	for {
 		iter := chunk.NewIterator4Chunk(chk)
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(chk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(chk))
 		if err != nil {
 			return err
 		}

executor/executor.go

@@ -175,6 +175,17 @@ type Executor interface {
 	newFirstChunk() *chunk.Chunk
 }
 
+// Next is a wrapper function on e.Next(), it handles some common codes.
+func Next(ctx context.Context, e Executor, req *chunk.RecordBatch) error {
+	select {
+	case <-ctx.Done():
+		req.Reset()
+		return errors.Trace(ctx.Err())
+	default:
+	}
+	return e.Next(ctx, req)
+}
+
 // CancelDDLJobsExec represents a cancel DDL jobs executor.
 type CancelDDLJobsExec struct {
 	baseExecutor
@@ -554,7 +565,7 @@ func (e *CheckIndexExec) Next(ctx context.Context, req *chunk.RecordBatch) error
 	}
 	chk := e.src.newFirstChunk()
 	for {
-		err := e.src.Next(ctx, chunk.NewRecordBatch(chk))
+		err := Next(ctx, e.src, chunk.NewRecordBatch(chk))
 		if err != nil {
 			return err
 		}
@@ -663,7 +674,7 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.RecordBatch) error
 	}
 
 	req.GrowAndReset(e.maxChunkSize)
-	err := e.children[0].Next(ctx, req)
+	err := Next(ctx, e.children[0], req)
 	if err != nil {
 		return err
 	}
@@ -723,7 +734,7 @@ func (e *LimitExec) Next(ctx context.Context, req *chunk.RecordBatch) error {
 	for !e.meetFirstBatch {
 		// transfer req's requiredRows to childResult and then adjust it in childResult
 		e.childResult = e.childResult.SetRequiredRows(req.RequiredRows(), e.maxChunkSize)
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(e.adjustRequiredRows(e.childResult)))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(e.adjustRequiredRows(e.childResult)))
 		if err != nil {
 			return err
 		}
@@ -748,7 +759,7 @@ func (e *LimitExec) Next(ctx context.Context, req *chunk.RecordBatch) error {
 		e.cursor += batchSize
 	}
 	e.adjustRequiredRows(req.Chunk)
-	err := e.children[0].Next(ctx, req)
+	err := Next(ctx, e.children[0], req)
 	if err != nil {
 		return err
 	}
@@ -818,7 +829,7 @@ func init() {
 		}
 		chk := exec.newFirstChunk()
 		for {
-			err = exec.Next(ctx, chunk.NewRecordBatch(chk))
+			err = Next(ctx, exec, chunk.NewRecordBatch(chk))
 			if err != nil {
 				return rows, err
 			}
@@ -935,7 +946,7 @@ func (e *SelectionExec) Next(ctx context.Context, req *chunk.RecordBatch) error
 			}
 			req.AppendRow(e.inputRow)
 		}
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(e.childResult))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(e.childResult))
 		if err != nil {
 			return err
 		}
@@ -967,7 +978,7 @@ func (e *SelectionExec) unBatchedNext(ctx context.Context, chk *chunk.Chunk) err
 				return nil
 			}
 		}
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(e.childResult))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(e.childResult))
 		if err != nil {
 			return err
 		}
@@ -1115,7 +1126,7 @@ func (e *MaxOneRowExec) Next(ctx context.Context, req *chunk.RecordBatch) error
 		return nil
 	}
 	e.evaluated = true
-	err := e.children[0].Next(ctx, req)
+	err := Next(ctx, e.children[0], req)
 	if err != nil {
 		return err
 	}
@@ -1130,7 +1141,7 @@ func (e *MaxOneRowExec) Next(ctx context.Context, req *chunk.RecordBatch) error
 	}
 
 	childChunk := e.children[0].newFirstChunk()
-	err = e.children[0].Next(ctx, chunk.NewRecordBatch(childChunk))
+	err = Next(ctx, e.children[0], chunk.NewRecordBatch(childChunk))
 	if err != nil {
 		return err
 	}
@@ -1241,7 +1252,7 @@ func (e *UnionExec) resultPuller(ctx context.Context, childID int) {
 			return
 		case result.chk = <-e.resourcePools[childID]:
 		}
-		result.err = e.children[childID].Next(ctx, chunk.NewRecordBatch(result.chk))
+		result.err = Next(ctx, e.children[childID], chunk.NewRecordBatch(result.chk))
 		if result.err == nil && result.chk.NumRows() == 0 {
 			return
 		}

executor/index_lookup_join.go

@@ -386,7 +386,7 @@ func (ow *outerWorker) buildTask(ctx context.Context) (*lookUpJoinTask, error) {
 
 	task.memTracker.Consume(task.outerResult.MemoryUsage())
 	for !task.outerResult.IsFull() {
-		err := ow.executor.Next(ctx, chunk.NewRecordBatch(ow.executorChk))
+		err := Next(ctx, ow.executor, chunk.NewRecordBatch(ow.executorChk))
 		if err != nil {
 			return task, err
 		}
@@ -586,7 +586,7 @@ func (iw *innerWorker) fetchInnerResults(ctx context.Context, task *lookUpJoinTa
 	innerResult.GetMemTracker().SetLabel(innerResultLabel)
 	innerResult.GetMemTracker().AttachTo(task.memTracker)
 	for {
-		err := innerExec.Next(ctx, chunk.NewRecordBatch(iw.executorChk))
+		err := Next(ctx, innerExec, chunk.NewRecordBatch(iw.executorChk))
 		if err != nil {
 			return err
 		}

executor/join.go

@@ -202,6 +202,7 @@ func (e *HashJoinExec) fetchOuterChunks(ctx context.Context) {
 		if e.finished.Load().(bool) {
 			return
 		}
+
 		var outerResource *outerChkResource
 		var ok bool
 		select {
@@ -217,7 +218,7 @@ func (e *HashJoinExec) fetchOuterChunks(ctx context.Context) {
 			required := int(atomic.LoadInt64(&e.requiredRows))
 			outerResult.SetRequiredRows(required, e.maxChunkSize)
 		}
-		err := e.outerExec.Next(ctx, chunk.NewRecordBatch(outerResult))
+		err := Next(ctx, e.outerExec, chunk.NewRecordBatch(outerResult))
 		if err != nil {
 			e.joinResultCh <- &hashjoinWorkerResult{
 				err: err,
@@ -244,6 +245,7 @@ func (e *HashJoinExec) fetchOuterChunks(ctx context.Context) {
 		if outerResult.NumRows() == 0 {
 			return
 		}
+
 		outerResource.dest <- outerResult
 	}
 }
@@ -276,8 +278,9 @@ func (e *HashJoinExec) fetchInnerRows(ctx context.Context) error {
 		if e.finished.Load().(bool) {
 			return nil
 		}
+
 		chk := e.children[e.innerIdx].newFirstChunk()
-		err = e.innerExec.Next(ctx, chunk.NewRecordBatch(chk))
+		err = Next(ctx, e.innerExec, chunk.NewRecordBatch(chk))
 		if err != nil || chk.NumRows() == 0 {
 			return err
 		}
@@ -512,6 +515,7 @@ func (e *HashJoinExec) Next(ctx context.Context, req *chunk.RecordBatch) (err er
 	if e.joinResultCh == nil {
 		return nil
 	}
+
 	result, ok := <-e.joinResultCh
 	if !ok {
 		return nil
@@ -642,7 +646,7 @@ func (e *NestedLoopApplyExec) fetchSelectedOuterRow(ctx context.Context, chk *ch
 	outerIter := chunk.NewIterator4Chunk(e.outerChunk)
 	for {
 		if e.outerChunkCursor >= e.outerChunk.NumRows() {
-			err := e.outerExec.Next(ctx, chunk.NewRecordBatch(e.outerChunk))
+			err := Next(ctx, e.outerExec, chunk.NewRecordBatch(e.outerChunk))
 			if err != nil {
 				return nil, err
 			}
@@ -679,7 +683,7 @@ func (e *NestedLoopApplyExec) fetchAllInners(ctx context.Context) error {
 	e.innerList.Reset()
 	innerIter := chunk.NewIterator4Chunk(e.innerChunk)
 	for {
-		err := e.innerExec.Next(ctx, chunk.NewRecordBatch(e.innerChunk))
+		err := Next(ctx, e.innerExec, chunk.NewRecordBatch(e.innerChunk))
 		if err != nil {
 			return err
 		}

executor/merge_join.go

@@ -142,7 +142,7 @@ func (t *mergeJoinInnerTable) nextRow() (chunk.Row, error) {
 		if t.curRow == t.curIter.End() {
 			t.reallocReaderResult()
 			oldMemUsage := t.curResult.MemoryUsage()
-			err := t.reader.Next(t.ctx, chunk.NewRecordBatch(t.curResult))
+			err := Next(t.ctx, t.reader, chunk.NewRecordBatch(t.curResult))
 			// error happens or no more data.
 			if err != nil || t.curResult.NumRows() == 0 {
 				t.curRow = t.curIter.End()
@@ -389,7 +389,7 @@ func (e *MergeJoinExec) fetchNextOuterRows(ctx context.Context, requiredRows int
 		e.outerTable.chk.SetRequiredRows(requiredRows, e.maxChunkSize)
 	}
 
-	err = e.outerTable.reader.Next(ctx, chunk.NewRecordBatch(e.outerTable.chk))
+	err = Next(ctx, e.outerTable.reader, chunk.NewRecordBatch(e.outerTable.chk))
 	if err != nil {
 		return err
 	}

executor/projection.go

@@ -179,7 +179,7 @@ func (e *ProjectionExec) isUnparallelExec() bool {
 func (e *ProjectionExec) unParallelExecute(ctx context.Context, chk *chunk.Chunk) error {
 	// transmit the requiredRows
 	e.childResult.SetRequiredRows(chk.RequiredRows(), e.maxChunkSize)
-	err := e.children[0].Next(ctx, chunk.NewRecordBatch(e.childResult))
+	err := Next(ctx, e.children[0], chunk.NewRecordBatch(e.childResult))
 	if err != nil {
 		return err
 	}
@@ -312,7 +312,7 @@ func (f *projectionInputFetcher) run(ctx context.Context) {
 
 		requiredRows := atomic.LoadInt64(&f.proj.parentReqRows)
 		input.chk.SetRequiredRows(int(requiredRows), f.proj.maxChunkSize)
-		err := f.child.Next(ctx, chunk.NewRecordBatch(input.chk))
+		err := Next(ctx, f.child, chunk.NewRecordBatch(input.chk))
 		if err != nil || input.chk.NumRows() == 0 {
 			output.done <- err
 			return

executor/sort.go

@@ -111,7 +111,7 @@ func (e *SortExec) fetchRowChunks(ctx context.Context) error {
 	e.rowChunks.GetMemTracker().SetLabel(rowChunksLabel)
 	for {
 		chk := e.children[0].newFirstChunk()
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(chk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(chk))
 		if err != nil {
 			return err
 		}
@@ -282,7 +282,7 @@ func (e *TopNExec) loadChunksUntilTotalLimit(ctx context.Context) error {
 		srcChk := e.children[0].newFirstChunk()
 		// adjust required rows by total limit
 		srcChk.SetRequiredRows(int(e.totalLimit-uint64(e.rowChunks.Len())), e.maxChunkSize)
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(srcChk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(srcChk))
 		if err != nil {
 			return err
 		}
@@ -307,7 +307,7 @@ func (e *TopNExec) executeTopN(ctx context.Context) error {
 	}
 	childRowChk := e.children[0].newFirstChunk()
 	for {
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(childRowChk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(childRowChk))
 		if err != nil {
 			return err
 		}

executor/union_scan.go

@@ -199,7 +199,7 @@ func (us *UnionScanExec) getSnapshotRow(ctx context.Context) ([]types.Datum, err
 	us.cursor4SnapshotRows = 0
 	us.snapshotRows = us.snapshotRows[:0]
 	for len(us.snapshotRows) == 0 {
-		err = us.children[0].Next(ctx, chunk.NewRecordBatch(us.snapshotChunkBuffer))
+		err = Next(ctx, us.children[0], chunk.NewRecordBatch(us.snapshotChunkBuffer))
 		if err != nil || us.snapshotChunkBuffer.NumRows() == 0 {
 			return nil, err
 		}

executor/update.go

@@ -181,7 +181,7 @@ func (e *UpdateExec) fetchChunkRows(ctx context.Context) error {
 	chk := e.children[0].newFirstChunk()
 	e.evalBuffer = chunk.MutRowFromTypes(fields)
 	for {
-		err := e.children[0].Next(ctx, chunk.NewRecordBatch(chk))
+		err := Next(ctx, e.children[0], chunk.NewRecordBatch(chk))
 		if err != nil {
 			return err
 		}

executor/window.go

@@ -131,7 +131,7 @@ func (e *WindowExec) fetchChildIfNecessary(ctx context.Context, chk *chunk.Chunk
 	}
 
 	childResult := e.children[0].newFirstChunk()
-	err = e.children[0].Next(ctx, &chunk.RecordBatch{Chunk: childResult})
+	err = Next(ctx, e.children[0], &chunk.RecordBatch{Chunk: childResult})
 	if err != nil {
 		return errors.Trace(err)
 	}

server/server.go

@@ -530,15 +530,8 @@ func (s *Server) Kill(connectionID uint64, query bool) {
 
 func killConn(conn *clientConn) {
 	conn.mu.RLock()
-	resultSets := conn.mu.resultSets
 	cancelFunc := conn.mu.cancelFunc
 	conn.mu.RUnlock()
-	for _, resultSet := range resultSets {
-		// resultSet.Close() is reentrant so it's safe to kill a same connID multiple times
-		if err := resultSet.Close(); err != nil {
-			logutil.Logger(context.Background()).Error("close result set error", zap.Uint32("connID", conn.connectionID), zap.Error(err))
-		}
-	}
 	if cancelFunc != nil {
 		cancelFunc()
 	}