disttask: fix potential panic in task executor (pingcap#50058)

ref pingcap#48795
nolouch · Jan 4, 2024 · cdc7572 · cdc7572
1 parent 494d9f0
commit cdc7572
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/pkg/disttask/framework/taskexecutor/manager.go b/pkg/disttask/framework/taskexecutor/manager.go
@@ -467,7 +467,7 @@ func (m *Manager) logErr(err error) {
 
 func (m *Manager) logErrAndPersist(err error, taskID int64, taskExecutor TaskExecutor) {
 	m.logErr(err)
-	if taskExecutor.IsRetryableError(err) {
+	if taskExecutor != nil && taskExecutor.IsRetryableError(err) {
 		logutil.Logger(m.logCtx).Error("met retryable err", zap.Error(err), zap.Stack("stack"))
 		return
 	}

diff --git a/pkg/disttask/framework/taskexecutor/manager_test.go b/pkg/disttask/framework/taskexecutor/manager_test.go
@@ -120,11 +120,30 @@ func TestOnRunnableTasks(t *testing.T) {
 	// no task
 	m.onRunnableTasks(nil)
 
+	// type not found
+	mockTaskTable.EXPECT().UpdateErrorToSubtask(m.ctx, id, taskID, gomock.Any())
+	m.onRunnableTask(task)
+
 	RegisterTaskType("type",
 		func(ctx context.Context, id string, task *proto.Task, taskTable TaskTable) TaskExecutor {
 			return mockInternalExecutor
 		})
 
+	// executor init failed non retryable
+	executorErr := errors.New("executor init failed")
+	mockInternalExecutor.EXPECT().Init(gomock.Any()).Return(executorErr)
+	mockInternalExecutor.EXPECT().IsRetryableError(executorErr).Return(false)
+	mockTaskTable.EXPECT().UpdateErrorToSubtask(m.ctx, id, taskID, executorErr)
+	m.onRunnableTask(task)
+	m.removeHandlingTask(taskID)
+	require.Equal(t, true, ctrl.Satisfied())
+
+	// executor init failed retryable
+	mockInternalExecutor.EXPECT().Init(gomock.Any()).Return(executorErr)
+	mockInternalExecutor.EXPECT().IsRetryableError(executorErr).Return(true)
+	m.onRunnableTask(task)
+	m.removeHandlingTask(taskID)
+
 	// get subtask failed
 	mockInternalExecutor.EXPECT().Init(gomock.Any()).Return(nil)
 	mockTaskTable.EXPECT().HasSubtasksInStates(m.ctx, id, taskID, proto.StepOne,