Merge branch 'master' of https://github.com/pingcap/tidb into pgo-rea…

…d-only

pkg/distsql/select_result.go

@@ -19,7 +19,6 @@ import (
 	"container/heap"
 	"context"
 	"fmt"
-	"maps"
 	"strconv"
 	"sync/atomic"
 	"time"
@@ -500,7 +499,6 @@ func (r *selectResult) updateCopRuntimeStats(ctx context.Context, copStats *copr
 	if r.rootPlanID <= 0 || r.ctx.RuntimeStatsColl == nil || (callee == "" && (copStats.ReqStats == nil || copStats.ReqStats.GetRPCStatsCount() == 0)) {
 		return
 	}
-
 	if copStats.ScanDetail.ProcessedKeys > 0 || copStats.TimeDetail.KvReadWallTime > 0 {
 		readKeys := copStats.ScanDetail.ProcessedKeys
 		readTime := copStats.TimeDetail.KvReadWallTime.Seconds()
@@ -655,7 +653,9 @@ func (s *selectResultRuntimeStats) mergeCopRuntimeStats(copStats *copr.CopRuntim
 		if s.backoffSleep == nil {
 			s.backoffSleep = make(map[string]time.Duration)
 		}
-		maps.Copy(s.backoffSleep, copStats.BackoffSleep)
+		for k, v := range copStats.BackoffSleep {
+			s.backoffSleep[k] += v
+		}
 	}
 	s.totalProcessTime += copStats.TimeDetail.ProcessTime
 	s.totalWaitTime += copStats.TimeDetail.WaitTime

pkg/distsql/select_result_test.go

@@ -17,6 +17,7 @@ package distsql
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
@@ -30,11 +31,15 @@ import (
 func TestUpdateCopRuntimeStats(t *testing.T) {
 	ctx := mock.NewContext()
 	ctx.GetSessionVars().StmtCtx = stmtctx.NewStmtCtx()
-	sr := selectResult{ctx: ctx.GetDistSQLCtx(), storeType: kv.TiKV}
+	sr := selectResult{ctx: ctx.GetDistSQLCtx(), storeType: kv.TiKV, stats: &selectResultRuntimeStats{}}
 	require.Nil(t, ctx.GetSessionVars().StmtCtx.RuntimeStatsColl)
 
 	sr.rootPlanID = 1234
-	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "a"}}, 0)
+	backOffSleep := make(map[string]time.Duration, 1)
+	backOffSleep["RegionMiss"] = time.Duration(100)
+	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "a", BackoffSleep: backOffSleep}}, 0)
+	// RuntimeStatsColl is nil, so the update doesn't take efffect
+	require.Equal(t, sr.stats.backoffSleep["RegionMiss"], time.Duration(0))
 
 	ctx.GetSessionVars().StmtCtx.RuntimeStatsColl = execdetails.NewRuntimeStatsColl(nil)
 	// refresh the ctx after assigning `RuntimeStatsColl`.
@@ -48,13 +53,17 @@ func TestUpdateCopRuntimeStats(t *testing.T) {
 
 	require.NotEqual(t, len(sr.copPlanIDs), len(sr.selectResp.GetExecutionSummaries()))
 
-	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "callee"}}, 0)
+	backOffSleep["RegionMiss"] = time.Duration(200)
+	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "callee", BackoffSleep: backOffSleep}}, 0)
 	require.False(t, ctx.GetSessionVars().StmtCtx.RuntimeStatsColl.ExistsCopStats(1234))
+	require.Equal(t, sr.stats.backoffSleep["RegionMiss"], time.Duration(200))
 
 	sr.copPlanIDs = []int{sr.rootPlanID}
 	require.NotNil(t, ctx.GetSessionVars().StmtCtx.RuntimeStatsColl)
 	require.Equal(t, len(sr.copPlanIDs), len(sr.selectResp.GetExecutionSummaries()))
 
-	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "callee"}}, 0)
+	backOffSleep["RegionMiss"] = time.Duration(300)
+	sr.updateCopRuntimeStats(context.Background(), &copr.CopRuntimeStats{CopExecDetails: execdetails.CopExecDetails{CalleeAddress: "callee", BackoffSleep: backOffSleep}}, 0)
 	require.Equal(t, "tikv_task:{time:1ns, loops:1}", ctx.GetSessionVars().StmtCtx.RuntimeStatsColl.GetCopStats(1234).String())
+	require.Equal(t, sr.stats.backoffSleep["RegionMiss"], time.Duration(500))
 }

pkg/disttask/framework/scheduler/balancer.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 	llog "github.com/pingcap/tidb/pkg/lightning/log"
 	"github.com/pingcap/tidb/pkg/util/intest"
@@ -116,6 +117,9 @@ func (b *balancer) doBalanceSubtasks(ctx context.Context, taskID int64, eligible
 	// managed nodes, subtasks of task might not be balanced.
 	adjustedNodes := filterNodesWithEnoughSlots(b.currUsedSlots, b.slotMgr.getCapacity(),
 		eligibleNodes, subtasks[0].Concurrency)
+	failpoint.Inject("mockNoEnoughSlots", func(_ failpoint.Value) {
+		adjustedNodes = []string{}
+	})
 	if len(adjustedNodes) == 0 {
 		// no node has enough slots to run the subtasks, skip balance and skip
 		// update used slots.

pkg/disttask/framework/taskexecutor/task_executor.go

@@ -122,7 +122,7 @@ func NewBaseTaskExecutor(ctx context.Context, task *proto.Task, param Param) *Ba
 //     `pending` state, to make sure subtasks can be balanced later when node scale out.
 //   - If current running subtask are scheduled away from this node, i.e. this node
 //     is taken as down, cancel running.
-func (e *BaseTaskExecutor) checkBalanceSubtask(ctx context.Context) {
+func (e *BaseTaskExecutor) checkBalanceSubtask(ctx context.Context, subtaskCtxCancel context.CancelFunc) {
 	ticker := time.NewTicker(checkBalanceSubtaskInterval)
 	defer ticker.Stop()
 	for {
@@ -143,7 +143,10 @@ func (e *BaseTaskExecutor) checkBalanceSubtask(ctx context.Context) {
 			e.logger.Info("subtask is scheduled away, cancel running",
 				zap.Int64("subtaskID", e.currSubtaskID.Load()))
 			// cancels runStep, but leave the subtask state unchanged.
-			e.cancelRunStepWith(nil)
+			if subtaskCtxCancel != nil {
+				subtaskCtxCancel()
+			}
+			failpoint.InjectCall("afterCancelSubtaskExec")
 			return
 		}
 
@@ -317,6 +320,12 @@ func (e *BaseTaskExecutor) Run() {
 				continue
 			}
 		}
+		if err := e.stepCtx.Err(); err != nil {
+			e.logger.Error("step executor context is done, the task should have been reverted",
+				zap.String("step", proto.Step2Str(task.Type, task.Step)),
+				zap.Error(err))
+			continue
+		}
 		err = e.runSubtask(subtask)
 		if err != nil {
 			// task executor keeps running its subtasks even though some subtask
@@ -418,23 +427,25 @@ func (e *BaseTaskExecutor) runSubtask(subtask *proto.Subtask) (resErr error) {
 	logTask := llog.BeginTask(logger, "run subtask")
 	subtaskErr := func() error {
 		e.currSubtaskID.Store(subtask.ID)
+		subtaskCtx, subtaskCtxCancel := context.WithCancel(e.stepCtx)
 
 		var wg util.WaitGroupWrapper
-		checkCtx, checkCancel := context.WithCancel(e.stepCtx)
+		checkCtx, checkCancel := context.WithCancel(subtaskCtx)
 		wg.RunWithLog(func() {
-			e.checkBalanceSubtask(checkCtx)
+			e.checkBalanceSubtask(checkCtx, subtaskCtxCancel)
 		})
 
 		if e.hasRealtimeSummary(e.stepExec) {
 			wg.RunWithLog(func() {
-				e.updateSubtaskSummaryLoop(checkCtx, e.stepCtx, e.stepExec)
+				e.updateSubtaskSummaryLoop(checkCtx, subtaskCtx, e.stepExec)
 			})
 		}
 		defer func() {
 			checkCancel()
 			wg.Wait()
+			subtaskCtxCancel()
 		}()
-		return e.stepExec.RunSubtask(e.stepCtx, subtask)
+		return e.stepExec.RunSubtask(subtaskCtx, subtask)
 	}()
 	failpoint.InjectCall("afterRunSubtask", e, &subtaskErr)
 	logTask.End2(zap.InfoLevel, subtaskErr)

pkg/disttask/framework/taskexecutor/task_executor_test.go

@@ -315,7 +315,7 @@ func TestTaskExecutorRun(t *testing.T) {
 		// mock for checkBalanceSubtask, returns empty subtask list
 		e.taskTable.EXPECT().GetSubtasksByExecIDAndStepAndStates(gomock.Any(), "id",
 			e.task1.ID, proto.StepOne, proto.SubtaskStateRunning).Return([]*proto.Subtask{}, nil)
-		// this subtask is scheduled awsy during running
+		// this subtask is scheduled away during running
 		e.taskTable.EXPECT().GetTaskByID(gomock.Any(), e.task1.ID).Return(e.task1, nil)
 		e.taskTable.EXPECT().GetFirstSubtaskInStates(gomock.Any(), "id", e.task1.ID, proto.StepOne,
 			unfinishedNormalSubtaskStates...).Return(e.pendingSubtask1, nil)
@@ -326,6 +326,7 @@ func TestTaskExecutorRun(t *testing.T) {
 			<-ctx.Done()
 			return ctx.Err()
 		})
+		e.taskExecExt.EXPECT().IsRetryableError(gomock.Any()).Return(true)
 		// keep running next subtask
 		nextSubtask := &proto.Subtask{SubtaskBase: proto.SubtaskBase{
 			ID: 2, Type: e.task1.Type, Step: proto.StepOne, State: proto.SubtaskStatePending, ExecID: "id"}}
@@ -889,18 +890,17 @@ func TestCheckBalanceSubtask(t *testing.T) {
 		// context canceled
 		canceledCtx, cancel := context.WithCancel(ctx)
 		cancel()
-		taskExecutor.checkBalanceSubtask(canceledCtx)
+		taskExecutor.checkBalanceSubtask(canceledCtx, nil)
 	})
 
 	t.Run("subtask scheduled away", func(t *testing.T) {
 		mockSubtaskTable.EXPECT().GetSubtasksByExecIDAndStepAndStates(gomock.Any(), "tidb1",
 			task.ID, task.Step, proto.SubtaskStateRunning).Return(nil, errors.New("error"))
 		mockSubtaskTable.EXPECT().GetSubtasksByExecIDAndStepAndStates(gomock.Any(), "tidb1",
 			task.ID, task.Step, proto.SubtaskStateRunning).Return([]*proto.Subtask{}, nil)
-		runCtx, cancelCause := context.WithCancelCause(ctx)
-		taskExecutor.mu.runtimeCancel = cancelCause
+		runCtx, cancel := context.WithCancel(ctx)
 		require.NoError(t, runCtx.Err())
-		taskExecutor.checkBalanceSubtask(ctx)
+		taskExecutor.checkBalanceSubtask(ctx, cancel)
 		require.ErrorIs(t, runCtx.Err(), context.Canceled)
 		require.True(t, ctrl.Satisfied())
 	})
@@ -913,7 +913,7 @@ func TestCheckBalanceSubtask(t *testing.T) {
 		mockExtension.EXPECT().IsIdempotent(subtasks[0]).Return(false)
 		mockSubtaskTable.EXPECT().UpdateSubtaskStateAndError(gomock.Any(), "tidb1",
 			subtasks[0].ID, proto.SubtaskStateFailed, ErrNonIdempotentSubtask).Return(nil)
-		taskExecutor.checkBalanceSubtask(ctx)
+		taskExecutor.checkBalanceSubtask(ctx, nil)
 		require.True(t, ctrl.Satisfied())
 
 		// if we failed to change state of non-idempotent subtask, will retry
@@ -930,7 +930,7 @@ func TestCheckBalanceSubtask(t *testing.T) {
 		mockExtension.EXPECT().IsIdempotent(subtasks[0]).Return(false)
 		mockSubtaskTable.EXPECT().UpdateSubtaskStateAndError(gomock.Any(), "tidb1",
 			subtasks[0].ID, proto.SubtaskStateFailed, ErrNonIdempotentSubtask).Return(nil)
-		taskExecutor.checkBalanceSubtask(ctx)
+		taskExecutor.checkBalanceSubtask(ctx, nil)
 		require.True(t, ctrl.Satisfied())
 	})
 
@@ -945,7 +945,7 @@ func TestCheckBalanceSubtask(t *testing.T) {
 		// used to break the loop
 		mockSubtaskTable.EXPECT().GetSubtasksByExecIDAndStepAndStates(gomock.Any(), "tidb1",
 			task.ID, task.Step, proto.SubtaskStateRunning).Return(nil, nil)
-		taskExecutor.checkBalanceSubtask(ctx)
+		taskExecutor.checkBalanceSubtask(ctx, nil)
 		require.True(t, ctrl.Satisfied())
 	})
 }

pkg/disttask/importinto/proto.go

@@ -110,7 +110,6 @@ type SharedVars struct {
 	TableImporter *importer.TableImporter
 	DataEngine    *backend.OpenedEngine
 	IndexEngine   *backend.OpenedEngine
-	Progress      *importer.Progress
 
 	mu       sync.Mutex
 	Checksum *verification.KVGroupChecksum
@@ -183,5 +182,4 @@ type Checksum struct {
 // This portion of the code may be implemented uniformly in the framework in the future.
 type Result struct {
 	LoadedRowCnt uint64
-	ColSizeMap   map[int64]int64
 }

pkg/disttask/importinto/scheduler.go

@@ -577,14 +577,9 @@ func updateResult(handle storage.TaskHandle, task *proto.Task, taskMeta *TaskMet
 		}
 		subtaskMetas = append(subtaskMetas, &subtaskMeta)
 	}
-	columnSizeMap := make(map[int64]int64)
 	for _, subtaskMeta := range subtaskMetas {
 		taskMeta.Result.LoadedRowCnt += subtaskMeta.Result.LoadedRowCnt
-		for key, val := range subtaskMeta.Result.ColSizeMap {
-			columnSizeMap[key] += val
-		}
 	}
-	taskMeta.Result.ColSizeMap = columnSizeMap
 
 	if globalSort {
 		taskMeta.Result.LoadedRowCnt, err = getLoadedRowCountOnGlobalSort(handle, task)
@@ -662,8 +657,7 @@ func (sch *ImportSchedulerExt) finishJob(ctx context.Context, logger *zap.Logger
 		func(ctx context.Context) (bool, error) {
 			return true, taskHandle.WithNewSession(func(se sessionctx.Context) error {
 				if err := importer.FlushTableStats(ctx, se, taskMeta.Plan.TableInfo.ID, &importer.JobImportResult{
-					Affected:   taskMeta.Result.LoadedRowCnt,
-					ColSizeMap: taskMeta.Result.ColSizeMap,
+					Affected: taskMeta.Result.LoadedRowCnt,
 				}); err != nil {
 					logger.Warn("flush table stats failed", zap.Error(err))
 				}

pkg/disttask/importinto/subtask_executor.go

@@ -69,7 +69,6 @@ func (e *importMinimalTaskExecutor) Run(ctx context.Context, dataWriter, indexWr
 			sharedVars.TableImporter,
 			sharedVars.DataEngine,
 			sharedVars.IndexEngine,
-			sharedVars.Progress,
 			logger,
 			checksum,
 		); err != nil {
@@ -82,7 +81,6 @@ func (e *importMinimalTaskExecutor) Run(ctx context.Context, dataWriter, indexWr
 			sharedVars.TableImporter,
 			dataWriter,
 			indexWriter,
-			sharedVars.Progress,
 			logger,
 			checksum,
 		); err != nil {

pkg/disttask/importinto/task_executor.go

@@ -154,7 +154,6 @@ func (s *importStepExecutor) RunSubtask(ctx context.Context, subtask *proto.Subt
 		TableImporter:    s.tableImporter,
 		DataEngine:       dataEngine,
 		IndexEngine:      indexEngine,
-		Progress:         importer.NewProgress(),
 		Checksum:         verification.NewKVGroupChecksumWithKeyspace(s.tableImporter.GetKeySpace()),
 		SortedDataMeta:   &external.SortedKVMeta{},
 		SortedIndexMetas: make(map[int64]*external.SortedKVMeta),
@@ -251,7 +250,6 @@ func (s *importStepExecutor) onFinished(ctx context.Context, subtask *proto.Subt
 	}
 	subtaskMeta.Result = Result{
 		LoadedRowCnt: dataKVCount,
-		ColSizeMap:   sharedVars.Progress.GetColSize(),
 	}
 	allocators := sharedVars.TableImporter.Allocators()
 	subtaskMeta.MaxIDs = map[autoid.AllocatorType]int64{

pkg/domain/domain.go

@@ -1902,7 +1902,7 @@ func (do *Domain) LoadPrivilegeLoop(sctx sessionctx.Context) error {
 	if err != nil {
 		return err
 	}
-	do.privHandle = privileges.NewHandle(sctx.GetRestrictedSQLExecutor())
+	do.privHandle = privileges.NewHandle(sctx.GetRestrictedSQLExecutor(), sctx.GetSessionVars().GlobalVarsAccessor)
 
 	var watchCh clientv3.WatchChan
 	duration := 5 * time.Minute

pkg/executor/delete.go

@@ -274,7 +274,7 @@ func (e *DeleteExec) removeRow(ctx sessionctx.Context, t table.Table, h kv.Handl
 		return err
 	}
 
-	err = t.RemoveRecord(ctx.GetTableCtx(), txn, h, data, posInfo.ExtraPartialRowOption)
+	err = t.RemoveRecord(ctx.GetTableCtx(), txn, h, data, posInfo.IndexesRowLayout)
 	if err != nil {
 		return err
 	}

pkg/executor/grant.go

@@ -155,6 +155,10 @@ func (e *GrantExec) Next(ctx context.Context, _ *chunk.Chunk) error {
 		return err
 	}
 
+	defaultAuthPlugin, err := e.Ctx().GetSessionVars().GlobalVarsAccessor.GetGlobalSysVar(variable.DefaultAuthPlugin)
+	if err != nil {
+		return err
+	}
 	// Check which user is not exist.
 	for _, user := range e.Users {
 		exists, err := userExists(ctx, e.Ctx(), user.User.Username, user.User.Hostname)
@@ -169,10 +173,7 @@ func (e *GrantExec) Next(ctx context.Context, _ *chunk.Chunk) error {
 			// It is required for compatibility with 5.7 but removed from 8.0
 			// since it results in a massive security issue:
 			// spelling errors will create users with no passwords.
-			authPlugin, err := e.Ctx().GetSessionVars().GlobalVarsAccessor.GetGlobalSysVar(variable.DefaultAuthPlugin)
-			if err != nil {
-				return err
-			}
+			authPlugin := defaultAuthPlugin
 			if user.AuthOpt != nil && user.AuthOpt.AuthPlugin != "" {
 				authPlugin = user.AuthOpt.AuthPlugin
 			}
@@ -181,7 +182,7 @@ func (e *GrantExec) Next(ctx context.Context, _ *chunk.Chunk) error {
 				return exeerrors.ErrPluginIsNotLoaded.GenWithStackByArgs(extErr.Error())
 			}
 			authPluginImpl := extensions.GetAuthPlugins()[authPlugin]
-			pwd, ok := encodePassword(user, authPluginImpl)
+			pwd, ok := encodePasswordWithPlugin(*user, authPluginImpl, defaultAuthPlugin)
 			if !ok {
 				return errors.Trace(exeerrors.ErrPasswordFormat)
 			}

pkg/executor/importer/BUILD.bazel

@@ -9,7 +9,6 @@ go_library(
         "job.go",
         "kv_encode.go",
         "precheck.go",
-        "progress.go",
         "table_import.go",
     ],
     importpath = "github.com/pingcap/tidb/pkg/executor/importer",

pkg/executor/importer/engine_process.go

@@ -30,7 +30,6 @@ func ProcessChunk(
 	chunk *checkpoints.ChunkCheckpoint,
 	tableImporter *TableImporter,
 	dataEngine, indexEngine *backend.OpenedEngine,
-	progress *Progress,
 	logger *zap.Logger,
 	groupChecksum *verification.KVGroupChecksum,
 ) error {
@@ -65,7 +64,7 @@ func ProcessChunk(
 		}
 	}()
 
-	return ProcessChunkWithWriter(ctx, chunk, tableImporter, dataWriter, indexWriter, progress, logger, groupChecksum)
+	return ProcessChunkWithWriter(ctx, chunk, tableImporter, dataWriter, indexWriter, logger, groupChecksum)
 }
 
 // ProcessChunkWithWriter processes a chunk, and write kv pairs to dataWriter and indexWriter.
@@ -74,7 +73,6 @@ func ProcessChunkWithWriter(
 	chunk *checkpoints.ChunkCheckpoint,
 	tableImporter *TableImporter,
 	dataWriter, indexWriter backend.EngineWriter,
-	progress *Progress,
 	logger *zap.Logger,
 	groupChecksum *verification.KVGroupChecksum,
 ) error {
@@ -116,6 +114,5 @@ func ProcessChunkWithWriter(
 	if err != nil {
 		return err
 	}
-	progress.AddColSize(encoder.GetColumnSize())
 	return nil
 }

pkg/executor/importer/import.go

@@ -1395,9 +1395,8 @@ func getDataSourceType(p *plannercore.ImportInto) DataSourceType {
 
 // JobImportResult is the result of the job import.
 type JobImportResult struct {
-	Affected   uint64
-	Warnings   []contextutil.SQLWarn
-	ColSizeMap variable.DeltaColsMap
+	Affected uint64
+	Warnings []contextutil.SQLWarn
 }
 
 // GetMsgFromBRError get msg from BR error.