pcm-coordinator/internal/logic/schedule/schedulecanceltasklogic.go

package schedule

import (
	"context"
	"errors"
	"fmt"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/utils/status"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"
	"strconv"

	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

type ScheduleCancelTaskLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewScheduleCancelTaskLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ScheduleCancelTaskLogic {
	return &ScheduleCancelTaskLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *ScheduleCancelTaskLogic) ScheduleCancelTask(req *types.CancelTaskReq) (resp *types.CancelTaskResp, err error) {
	task, err := l.svcCtx.Scheduler.AiService.Storage.GetTaskById(req.TaskId)
	if err != nil {
		return nil, err
	}

	if task == nil {
		return nil, errors.New("failed to cancel task, task not found")
	}

	// find ai tasks
	aitasks, err := l.svcCtx.Scheduler.AiStorages.GetAiTaskListById(task.Id)
	if err != nil {
		return nil, err
	}

	if len(aitasks) == 0 {
		return nil, errors.New("failed to cancel task, ai sub tasks have not been created")
	}

	// update status
	status.UpdateAiTask(l.svcCtx, aitasks...)

	t := aitasks[0]

	if t.Status != constants.Running {
		return nil, fmt.Errorf("failed to cancel task, ai sub tasks is %s", t.Status)
	}

	// assume a task has only one sub ai task
	err = l.svcCtx.Scheduler.AiService.AiExecutorAdapterMap[strconv.FormatInt(t.AdapterId, 10)][strconv.FormatInt(t.ClusterId, 10)].Stop(l.ctx, t.JobId)
	if err != nil {
		return nil, err
	}

	t.Status = constants.Cancelled
	err = l.svcCtx.Scheduler.AiService.Storage.UpdateAiTask(t)
	if err != nil {
		return nil, err
	}

	task.Status = constants.Cancelled
	err = l.svcCtx.Scheduler.AiService.Storage.UpdateTaskByModel(task)
	if err != nil {
		return nil, err
	}

	return resp, nil
}