baize/convenient/domain/task_runner/task_runner.go

package task_runner

import (
	"encoding/json"
	"git.sxidc.com/go-framework/baize/convenient/domain/task_runner/task"
	"git.sxidc.com/go-framework/baize/framework/core/domain"
	"git.sxidc.com/go-framework/baize/framework/core/domain/entity"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/database"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/database/sql"
	"git.sxidc.com/go-framework/baize/framework/core/infrastructure/logger"
	"git.sxidc.com/go-tools/utils/strutils"
	"github.com/pkg/errors"
	"sync"
)

var runnerRegister = new(sync.Map)

type RunTaskParams struct {
	Group      string
	OperatorID string
	Context    map[string]any
	DBSchema   string
	DBExecutor database.Executor
}

func (params *RunTaskParams) check() error {
	if strutils.IsStringEmpty(params.Group) {
		return errors.New("没有传递任务组")
	}

	if strutils.IsStringEmpty(params.OperatorID) {
		return errors.New("没有传递操作人ID")
	}

	if strutils.IsStringEmpty(params.DBSchema) {
		return errors.New("没有传递数据库schema")
	}

	if params.DBExecutor == nil {
		return errors.New("没有传递数据库执行器")
	}

	return nil
}

func RunTask(runner Runner, params *RunTaskParams) (string, error) {
	err := params.check()
	if err != nil {
		return "", err
	}

	ctxJsonBytes, err := json.Marshal(params.Context)
	if err != nil {
		return "", err
	}

	ctxJsonStr := string(ctxJsonBytes)

	// 创建任务，任务状态为已创建
	taskEntity := &task.Entity{
		Group:        params.Group,
		Context:      ctxJsonStr,
		CreateUserID: params.OperatorID,
	}

	err = createTaskDB(taskEntity, params.DBSchema, params.DBExecutor)
	if err != nil {
		return "", err
	}

	loaded, _ := runnerRegister.LoadOrStore(params.Group, runner)
	loadedRunner := loaded.(Runner)

	go runTask(taskEntity, loadedRunner, params.DBSchema, params.DBExecutor, func(ctx map[string]any, runner Runner) error {
		return runner.Run(ctx)
	})

	return taskEntity.ID, nil
}

type RestartTaskParams struct {
	Group      string
	DBSchema   string
	DBExecutor database.Executor
}

func (params *RestartTaskParams) check() error {
	if strutils.IsStringEmpty(params.Group) {
		return errors.New("没有传递任务组")
	}

	if strutils.IsStringEmpty(params.DBSchema) {
		return errors.New("没有传递数据库schema")
	}

	if params.DBExecutor == nil {
		return errors.New("没有传递数据库执行器")
	}

	return nil
}

func RestartTask(runner Runner, params *RestartTaskParams) error {
	err := params.check()
	if err != nil {
		return err
	}

	runningResults, _, err := database.Query(params.DBExecutor, &sql.QueryExecuteParams{
		TableName: domain.TableName(params.DBSchema, &task.Entity{}),
		Conditions: sql.NewConditions().Equal(task.ColumnGroup, params.Group).
			Equal(task.ColumnStatus, task.StatusCodeRunning),
		PageNo:   0,
		PageSize: 0,
	})
	if err != nil {
		return err
	}

	taskEntities := make([]task.Entity, 0)
	err = sql.ParseSqlResult(runningResults, &taskEntities)
	if err != nil {
		return err
	}

	for _, taskEntity := range taskEntities {
		loaded, _ := runnerRegister.LoadOrStore(params.Group, runner)
		loadedRunner := loaded.(Runner)

		go runTask(&taskEntity, loadedRunner, params.DBSchema, params.DBExecutor, func(ctx map[string]any, runner Runner) error {
			return runner.Restart(ctx)
		})
	}

	return nil
}

type StopTaskParams struct {
	ID         string
	DBSchema   string
	DBExecutor database.Executor
}

func (params *StopTaskParams) check() error {
	if strutils.IsStringEmpty(params.ID) {
		return errors.New("没有传递任务ID")
	}

	if strutils.IsStringEmpty(params.DBSchema) {
		return errors.New("没有传递数据库schema")
	}

	if params.DBExecutor == nil {
		return errors.New("没有传递数据库执行器")
	}

	return nil
}
func StopTask(params *StopTaskParams) error {
	err := params.check()
	if err != nil {
		return err
	}

	result, err := database.QueryOne(params.DBExecutor, &sql.QueryOneExecuteParams{
		TableName: domain.TableName(params.DBSchema, &task.Entity{}),
		Conditions: sql.NewConditions().Equal(entity.ColumnID, params.ID).
			Equal(task.ColumnStatus, task.StatusCodeRunning),
	})
	if err != nil {
		if database.IsErrorDBRecordNotExist(err) {
			return errors.New("任务不存在")
		}

		return err
	}

	taskEntity := new(task.Entity)
	err = sql.ParseSqlResult(result, taskEntity)
	if err != nil {
		return err
	}

	ctx, err := taskEntity.GetMapContext()
	if err != nil {
		updateErr := updateTaskStatusErrorDB(taskEntity, err.Error(), params.DBSchema, params.DBExecutor)
		if updateErr != nil {
			return updateErr
		}

		return err
	}

	loaded, ok := runnerRegister.Load(taskEntity.Group)
	if !ok {
		return errors.New("没有找到任务对应的执行器")
	}

	loadedRunner := loaded.(Runner)

	err = database.Transaction(params.DBExecutor, func(tx database.Executor) error {
		err := updateTaskStatusStopDBTx(taskEntity, params.DBSchema, tx)
		if err != nil {
			return err
		}

		err = loadedRunner.Stop(ctx)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}

func createTaskDB(taskEntity *task.Entity, dbSchema string, dbExecutor database.Executor) error {
	taskEntity.SetStatusCreated()

	err := domain.CheckWhen(taskEntity, taskEntity.GetFieldMap(), "run")
	if err != nil {
		return err
	}

	err = database.InsertEntity(dbExecutor, domain.TableName(dbSchema, taskEntity), taskEntity)
	if err != nil {
		return err
	}

	return nil
}

func updateTaskStatusRunningDB(taskEntity *task.Entity, dbSchema string, dbExecutor database.Executor) error {
	taskEntity.SetStatusRunning()

	err := database.UpdateEntity(dbExecutor, domain.TableName(dbSchema, taskEntity), taskEntity)
	if err != nil {
		return err
	}

	return nil
}

func updateTaskStatusCompleteDB(taskEntity *task.Entity, dbSchema string, dbExecutor database.Executor) error {
	taskEntity.SetStatusCompleted()

	err := database.UpdateEntity(dbExecutor, domain.TableName(dbSchema, taskEntity), taskEntity)
	if err != nil {
		return err
	}

	return nil
}

func updateTaskStatusStopDBTx(taskEntity *task.Entity, dbSchema string, tx database.Executor) error {
	taskEntity.SetStatusStop()

	err := database.UpdateEntity(tx, domain.TableName(dbSchema, taskEntity), taskEntity)
	if err != nil {
		return err
	}

	return nil
}

func updateTaskStatusErrorDB(taskEntity *task.Entity, errMsg string, dbSchema string, dbExecutor database.Executor) error {
	taskEntity.SetStatusError(errMsg)

	err := database.UpdateEntity(dbExecutor, domain.TableName(dbSchema, taskEntity), taskEntity)
	if err != nil {
		return err
	}

	return nil
}

func runTask(taskEntity *task.Entity, runner Runner, dbSchema string, dbExecutor database.Executor, executeFunc func(ctx map[string]any, runner Runner) error) {
	ctx, err := taskEntity.GetMapContext()
	if err != nil {
		err = updateTaskStatusErrorDB(taskEntity, err.Error(), dbSchema, dbExecutor)
		if err != nil {
			logger.GetInstance().Error("解析任务上下文失败 ", err)
			return
		}

		return
	}

	err = updateTaskStatusRunningDB(taskEntity, dbSchema, dbExecutor)
	if err != nil {
		logger.GetInstance().Error("更新任务运行状态失败 ", err)
		return
	}

	// 执行任务
	err = executeFunc(ctx, runner)
	if err != nil {
		err = updateTaskStatusErrorDB(taskEntity, err.Error(), dbSchema, dbExecutor)
		if err != nil {
			logger.GetInstance().Error("更新任务错误状态失败 ", err)
			return
		}

		return
	}

	err = updateTaskStatusCompleteDB(taskEntity, dbSchema, dbExecutor)
	if err != nil {
		logger.GetInstance().Error("更新任务完成状态失败 ", err)
		return
	}
}