update

Author: CbcWestwolf

pkg/disttask/framework/scheduler/balancer.go

@@ -74,7 +74,7 @@ func (b *balancer) balanceLoop(ctx context.Context, sm *Manager) {
 func (b *balancer) balance(ctx context.Context, sm *Manager) {
 	// we will use currUsedSlots to calculate adjusted eligible nodes during balance,
 	// it's initial value depends on the managed nodes, to have a consistent view,
-	// DO NOT call getManagedNodes twice during 1 balance.
+	// DO NOT call getNodes twice during 1 balance.
 	managedNodes := b.nodeMgr.getNodes()
 	b.currUsedSlots = make(map[string]int, len(managedNodes))
 	for _, n := range managedNodes {

pkg/disttask/framework/taskexecutor/task_executor.go

@@ -447,10 +447,12 @@ func (e *BaseTaskExecutor) runSubtask(subtask *proto.Subtask) (resErr error) {
 		defer func() {
 			checkCancel()
 			wg.Wait()
-			subtaskCtxCancel()
 		}()
 		return e.stepExec.RunSubtask(subtaskCtx, subtask)
 	}()
+	defer func() {
+		subtaskCtxCancel()
+	}()
 	failpoint.InjectCall("afterRunSubtask", e, &subtaskErr)
 	logTask.End2(zap.InfoLevel, subtaskErr)
 
@@ -463,7 +465,7 @@ func (e *BaseTaskExecutor) runSubtask(subtask *proto.Subtask) (resErr error) {
 		return subtaskErr
 	}
 
-	err := e.finishSubtask(e.stepCtx, subtask)
+	err := e.finishSubtask(subtaskCtx, subtask)
 	failpoint.InjectCall("syncAfterSubtaskFinish")
 	return err
 }
@@ -570,7 +572,7 @@ func (e *BaseTaskExecutor) markSubTaskCanceledOrFailed(ctx context.Context, subt
 			return e.updateSubtaskStateAndErrorImpl(e.ctx, subtask.ExecID, subtask.ID, proto.SubtaskStateCanceled, nil)
 		}
 
-		e.logger.Info("meet context canceled for gracefully shutdown")
+		e.logger.Info("meet context canceled for gracefully shutdown", zap.Error(ctx.Err()))
 	} else if e.IsRetryableError(stErr) {
 		e.logger.Warn("meet retryable error", zap.Error(stErr))
 	} else {