crawlab/core/controllers/task.go

package controllers

import (
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/crawlab-team/crawlab/core/constants"
	"github.com/crawlab-team/crawlab/core/interfaces"
	"github.com/crawlab-team/crawlab/core/models/models"
	"github.com/crawlab-team/crawlab/core/models/service"
	mongo3 "github.com/crawlab-team/crawlab/core/mongo"
	"github.com/crawlab-team/crawlab/core/spider/admin"
	"github.com/crawlab-team/crawlab/core/task/log"
	"github.com/crawlab-team/crawlab/core/task/scheduler"
	"github.com/crawlab-team/crawlab/core/utils"
	"github.com/gin-gonic/gin"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	mongo2 "go.mongodb.org/mongo-driver/mongo"
)

type GetTaskByIdParams struct {
	Id string `path:"id"`
}

func GetTaskById(_ *gin.Context, params *GetTaskByIdParams) (response *Response[models.Task], err error) {
	// id
	id, err := primitive.ObjectIDFromHex(params.Id)
	if err != nil {
		return GetErrorResponse[models.Task](err)
	}

	// task
	t, err := service.NewModelService[models.Task]().GetById(id)
	if errors.Is(err, mongo2.ErrNoDocuments) {
		return GetErrorResponse[models.Task](err)
	}
	if err != nil {
		return GetErrorResponse[models.Task](err)
	}

	// skip if task status is pending
	if t.Status == constants.TaskStatusPending {
		return GetDataResponse(*t)
	}

	// spider
	if !t.SpiderId.IsZero() {
		t.Spider, _ = service.NewModelService[models.Spider]().GetById(t.SpiderId)
	}

	// schedule
	if !t.ScheduleId.IsZero() {
		t.Schedule, _ = service.NewModelService[models.Schedule]().GetById(t.ScheduleId)
	}

	// node
	if !t.NodeId.IsZero() {
		t.Node, _ = service.NewModelService[models.Node]().GetById(t.NodeId)
	}

	// task stat
	t.Stat, _ = service.NewModelService[models.TaskStat]().GetById(id)

	return GetDataResponse(*t)
}

type GetTaskListParams struct {
	*GetListParams
	Stats bool `query:"stats"`
}

func GetTaskList(c *gin.Context, params *GetTaskListParams) (response *ListResponse[models.Task], err error) {
	if params.Stats {
		return NewController[models.Task]().GetList(c, params.GetListParams)
	}

	// get query
	query, err := GetFilterQueryFromListParams(params.GetListParams)
	if err != nil {
		return GetErrorListResponse[models.Task](err)
	}

	sort, err := GetSortOptionFromString(params.GetListParams.Sort)
	if err != nil {
		return GetErrorListResponse[models.Task](err)
	}

	// get tasks
	tasks, err := service.NewModelService[models.Task]().GetMany(query, &mongo3.FindOptions{
		Sort:  sort,
		Skip:  params.Size * (params.Page - 1),
		Limit: params.Size,
	})
	if err != nil {
		if errors.Is(err, mongo2.ErrNoDocuments) {
			return GetErrorListResponse[models.Task](err)
		}
		return GetErrorListResponse[models.Task](err)
	}

	// check empty list
	if len(tasks) == 0 {
		return GetListResponse[models.Task](nil, 0)
	}

	// ids
	var taskIds []primitive.ObjectID
	var spiderIds []primitive.ObjectID
	for _, t := range tasks {
		taskIds = append(taskIds, t.Id)
		spiderIds = append(spiderIds, t.SpiderId)
	}

	// total count
	total, err := service.NewModelService[models.Task]().Count(query)
	if err != nil {
		return GetErrorListResponse[models.Task](err)
	}

	// stat list
	stats, err := service.NewModelService[models.TaskStat]().GetMany(bson.M{
		"_id": bson.M{
			"$in": taskIds,
		},
	}, nil)
	if err != nil {
		return GetErrorListResponse[models.Task](err)
	}

	// cache stat list to dict
	statsDict := map[primitive.ObjectID]models.TaskStat{}
	for _, s := range stats {
		statsDict[s.Id] = s
	}

	// spider list
	spiders, err := service.NewModelService[models.Spider]().GetMany(bson.M{
		"_id": bson.M{
			"$in": spiderIds,
		},
	}, nil)
	if err != nil {
		return GetErrorListResponse[models.Task](err)
	}

	// cache spider list to dict
	spiderDict := map[primitive.ObjectID]models.Spider{}
	for _, s := range spiders {
		spiderDict[s.Id] = s
	}

	// iterate list again
	for i, t := range tasks {
		// task stat
		ts, ok := statsDict[t.Id]
		if ok {
			tasks[i].Stat = &ts
		}

		// spider
		s, ok := spiderDict[t.SpiderId]
		if ok {
			tasks[i].Spider = &s
		}
	}

	return GetListResponse(tasks, total)
}

type DeleteTaskByIdParams struct {
	Id string `path:"id"`
}

func DeleteTaskById(_ *gin.Context, params *DeleteTaskByIdParams) (response *VoidResponse, err error) {
	id, err := primitive.ObjectIDFromHex(params.Id)
	if err != nil {
		return GetErrorVoidResponse(err)
	}

	// delete in db
	if err := mongo3.RunTransaction(func(context mongo2.SessionContext) (err error) {
		// delete task
		_, err = service.NewModelService[models.Task]().GetById(id)
		if err != nil {
			return err
		}
		err = service.NewModelService[models.Task]().DeleteById(id)
		if err != nil {
			return err
		}

		// delete task stat
		_, err = service.NewModelService[models.TaskStat]().GetById(id)
		if err != nil {
			logger.Warnf("delete task stat error: %s", err.Error())
			return nil
		}
		err = service.NewModelService[models.TaskStat]().DeleteById(id)
		if err != nil {
			logger.Warnf("delete task stat error: %s", err.Error())
			return nil
		}

		return nil
	}); err != nil {
		return GetErrorVoidResponse(err)
	}

	// delete task logs
	logPath := filepath.Join(utils.GetTaskLogPath(), id.Hex())
	if err := os.RemoveAll(logPath); err != nil {
		logger.Warnf("failed to remove task log directory: %s", logPath)
	}

	return GetVoidResponse()
}

type DeleteTaskListParams struct {
	Ids []string `json:"ids"`
}

func DeleteList(_ *gin.Context, params *DeleteTaskListParams) (response *VoidResponse, err error) {
	var ids []primitive.ObjectID
	for _, id := range params.Ids {
		id, err := primitive.ObjectIDFromHex(id)
		if err != nil {
			return GetErrorVoidResponse(err)
		}
		ids = append(ids, id)
	}

	if err := mongo3.RunTransaction(func(context mongo2.SessionContext) error {
		// delete tasks
		if err := service.NewModelService[models.Task]().DeleteMany(bson.M{
			"_id": bson.M{
				"$in": ids,
			},
		}); err != nil {
			return err
		}

		// delete task stats
		if err := service.NewModelService[models.Task]().DeleteMany(bson.M{
			"_id": bson.M{
				"$in": ids,
			},
		}); err != nil {
			logger.Warnf("delete task stat error: %s", err.Error())
			return nil
		}

		return nil
	}); err != nil {
		return GetErrorVoidResponse(err)
	}

	// delete tasks logs
	wg := sync.WaitGroup{}
	wg.Add(len(ids))
	for _, id := range ids {
		go func(taskId primitive.ObjectID) {
			// delete task logs
			logPath := filepath.Join(utils.GetTaskLogPath(), taskId.Hex())
			if err := os.RemoveAll(logPath); err != nil {
				logger.Warnf("failed to remove task log directory: %s", logPath)
			}
			wg.Done()
		}(id)
	}
	wg.Wait()

	return GetVoidResponse()
}

type PostTaskRunParams struct {
	SpiderId string   `json:"spider_id" validate:"required"`
	Mode     string   `json:"mode"`
	NodeIds  []string `json:"node_ids"`
	Cmd      string   `json:"cmd"`
	Param    string   `json:"param"`
	Priority int      `json:"priority"`
}

func PostTaskRun(c *gin.Context, params *PostTaskRunParams) (response *Response[[]primitive.ObjectID], err error) {
	spiderId, err := primitive.ObjectIDFromHex(params.SpiderId)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	var nodeIds []primitive.ObjectID
	if params.NodeIds != nil {
		for _, nodeId := range params.NodeIds {
			nodeId, err := primitive.ObjectIDFromHex(nodeId)
			if err != nil {
				return GetErrorResponse[[]primitive.ObjectID](err)
			}
			nodeIds = append(nodeIds, nodeId)
		}
	}

	// spider
	s, err := service.NewModelService[models.Spider]().GetById(spiderId)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	// options
	opts := &interfaces.SpiderRunOptions{
		Mode:     params.Mode,
		NodeIds:  nodeIds,
		Cmd:      params.Cmd,
		Param:    params.Param,
		Priority: params.Priority,
	}

	// user
	if u := GetUserFromContext(c); u != nil {
		opts.UserId = u.Id
	}

	// run
	adminSvc := admin.GetSpiderAdminService()
	taskIds, err := adminSvc.Schedule(s.Id, opts)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	return GetDataResponse(taskIds)
}

type PostTaskRestartParams struct {
	Id string `path:"id"`
}

func PostTaskRestart(c *gin.Context, params *PostTaskRestartParams) (response *Response[[]primitive.ObjectID], err error) {
	// id
	id, err := primitive.ObjectIDFromHex(params.Id)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	// task
	t, err := service.NewModelService[models.Task]().GetById(id)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	// options
	opts := &interfaces.SpiderRunOptions{
		Mode:     t.Mode,
		NodeIds:  t.NodeIds,
		Cmd:      t.Cmd,
		Param:    t.Param,
		Priority: t.Priority,
	}

	// normalize options for tasks with assigned node
	if !t.NodeId.IsZero() {
		opts.NodeIds = []primitive.ObjectID{t.NodeId}
		opts.Mode = constants.RunTypeSelectedNodes
	}

	// user
	if u := GetUserFromContext(c); u != nil {
		opts.UserId = u.Id
	}

	// run
	adminSvc := admin.GetSpiderAdminService()
	taskIds, err := adminSvc.Schedule(t.SpiderId, opts)
	if err != nil {
		return GetErrorResponse[[]primitive.ObjectID](err)
	}

	return GetDataResponse(taskIds)
}

type PostTaskCancelParams struct {
	Id    string `path:"id"`
	Force bool   `json:"force,omitempty"`
}

func PostTaskCancel(c *gin.Context, params *PostTaskCancelParams) (response *VoidResponse, err error) {
	// id
	id, err := primitive.ObjectIDFromHex(params.Id)
	if err != nil {
		return GetErrorVoidResponse(err)
	}

	// task
	t, err := service.NewModelService[models.Task]().GetById(id)
	if err != nil {
		return GetErrorVoidResponse(err)
	}

	// validate
	if !utils.IsCancellable(t.Status) {
		return GetErrorVoidResponse(errors.New("task is not cancellable"))
	}

	u := GetUserFromContext(c)

	// cancel
	schedulerSvc := scheduler.GetTaskSchedulerService()
	err = schedulerSvc.Cancel(id, u.Id, params.Force)
	if err != nil {
		return GetErrorVoidResponse(err)
	}

	return GetVoidResponse()
}

type GetTaskLogsParams struct {
	Id   string `path:"id"`
	Page int    `query:"page"`
	Size int    `query:"size"`
}

func GetTaskLogs(_ *gin.Context, params *GetTaskLogsParams) (response *ListResponse[string], err error) {
	// id
	id, err := primitive.ObjectIDFromHex(params.Id)
	if err != nil {
		return GetErrorListResponse[string](err)
	}

	// logs
	logDriver := log.GetFileLogDriver()
	logs, err := logDriver.Find(id.Hex(), "", (params.Page-1)*params.Size, params.Size)
	if err != nil {
		if strings.HasSuffix(err.Error(), "Status:404 Not Found") {
			return GetListResponse[string](nil, 0)
		}
		return GetErrorListResponse[string](err)
	}
	total, err := logDriver.Count(id.Hex(), "")
	if err != nil {
		return GetErrorListResponse[string](err)
	}

	return GetListResponse(logs, total)
}