remote-task-excutor-cli/pkg/service/task/task_service.go

package task

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"remote-task-excutor-cli/pkg/client"
	"remote-task-excutor-cli/pkg/models"
	"remote-task-excutor-cli/pkg/service"
	"remote-task-excutor-cli/pkg/service/common"
	"remote-task-excutor-cli/pkg/service/updown"
	"time"

	"github.com/mholt/archiver/v3"
)

type taskService struct {
	httpClient    *client.HTTPClient
	taskSubmitter *common.TaskSubmitter
	updownService service.UpDownService
}

func NewTaskService(baseURL string, timeout time.Duration) service.TaskService {
	httpClient := client.NewHTTPClient(baseURL, timeout)
	return &taskService{
		httpClient:    httpClient,
		taskSubmitter: common.NewTaskSubmitter(httpClient, common.TrainingTaskType),
		updownService: updown.NewUpDownService(baseURL, timeout),
	}
}

func (s *taskService) SubmitTask(ctx context.Context, authData *models.AuthData, config *models.RunConfig,
	clusterID string, bindResultSet *models.BindResultSet) (string, error) {
	jobSetID, _, err := s.taskSubmitter.SubmitTask(ctx, authData, config, clusterID, bindResultSet)
	return jobSetID, err
}

func (s *taskService) GetTaskStatus(ctx context.Context, authData *models.AuthData, jobSetID string) (*models.TaskDetailResponse, error) {
	s.httpClient.SetHeader("Authorization", "Bearer "+authData.Token)
	params := map[string]string{
		models.JobSetIDKey:   jobSetID,
		models.LocalJobIDKey: models.MainTaskID,
	}

	for {
		resp, err := s.httpClient.Get("/jsm/v2/jobs/details", params)
		if err != nil {
			fmt.Println("Get task status failed: ", err)
			return nil, err
		}

		var taskStatus models.TaskDetailResponse
		if err := json.Unmarshal(resp, &taskStatus); err != nil {
			fmt.Println("Get task status response unmarshal failed: ", err)
			return nil, err
		}

		if taskStatus.Code != models.ResponseOK {
			fmt.Println("Get task status failed: ", taskStatus.Code)
			return nil, fmt.Errorf("get task status failed: code: %s, message: %s", taskStatus.Code, taskStatus.Message)
		}
		fmt.Println("Get task status result: ", string(resp))
		if len(taskStatus.Data.SubTaskInfos) == 0 {
			time.Sleep(5 * time.Second)
			continue
		}
		return &taskStatus, nil
	}
}

func (s *taskService) GetTaskResult(ctx context.Context, authData *models.AuthData, path string, jobSetID string) error {
	result, err := s.getTaskResultDetail(ctx, authData, jobSetID)
	if err != nil {
		return err
	}

	filename, err := s.updownService.DownLoad(ctx, authData, path, result.PcmJobData.ResultFiles[0].Objects[0].PackageID)
	if err != nil {
		return err
	}

	if err := archiver.Unarchive(filename, path); err != nil {
		fmt.Println("Unarchive task result failed: ", err)
		return err
	}

	// 删除filename文件
	if err := os.Remove(filename); err != nil {
		fmt.Println("Remove task result failed: ", err)
		return err
	}

	return nil
}

func (s *taskService) getTaskResultDetail(ctx context.Context, authData *models.AuthData,
	jobSetID string) (*models.TaskResultDetailData, error) {
	s.httpClient.SetHeader("Authorization", "Bearer "+authData.Token)
	queryTaskResultReq := models.TaskQueryRequest{
		JobSetID:   jobSetID,
		LocalJobID: models.MainTaskID,
	}
	//todo 最多查询10分钟，10分钟没有结果就放弃查询，不管了
	maxAttemps := 3
	attemps := 0
	for {
		resp, err := s.httpClient.PostJSON("/jsm/v2/jobs/results", queryTaskResultReq)
		fmt.Println("get result details resp:", string(resp))
		if err != nil {
			fmt.Println("Get task result details failed: ", err)
			return nil, err
		}

		var taskResultDetails models.TaskResultDetailResponse
		if err := json.Unmarshal(resp, &taskResultDetails); err != nil {
			fmt.Println("Get task result details response unmarshal failed: ", err)
			return nil, err
		}

		if taskResultDetails.Code != models.ResponseOK {
			if taskResultDetails.Message == "query pcm job: record not found" {
				time.Sleep(5 * time.Second)
				continue
			} else {
				if attemps > maxAttemps {
					fmt.Println("Get task result details failed: ", taskResultDetails.Code)
					return nil, fmt.Errorf("get task result details failed: %s", taskResultDetails.Code)
				}
				attemps++
				time.Sleep(5 * time.Second)
			}
		}

		if taskResultDetails.Data.PcmJobData.Status == "failed" { // && taskResultDetails.Data.PcmJobData.ErrorMsg == "data return job id is empty" {
			time.Sleep(5 * time.Second)
			continue
		}

		return &taskResultDetails.Data, nil
	}

}