content: package com.aizuda.snailjob.server.job.task.support.dispatch;
  
  import com.aizuda.snailjob.server.common.dto.InstanceLiveInfo;
  import com.aizuda.snailjob.server.common.handler.InstanceManager;
  import org.apache.pekko.actor.AbstractActor;
  import cn.hutool.core.collection.CollUtil;
  import cn.hutool.core.lang.Assert;
  import cn.hutool.core.util.StrUtil;
  import com.aizuda.snailjob.common.core.constant.SystemConstants;
  import com.aizuda.snailjob.common.core.context.SnailSpringContext;
  import com.aizuda.snailjob.common.core.enums.*;
  import com.aizuda.snailjob.common.core.util.JsonUtil;
  import com.aizuda.snailjob.common.log.SnailJobLog;
  import com.aizuda.snailjob.server.common.pekko.ActorGenerator;
  import com.aizuda.snailjob.server.common.enums.JobTaskExecutorSceneEnum;
  import com.aizuda.snailjob.server.common.exception.SnailJobServerException;
  import com.aizuda.snailjob.server.common.util.DateUtils;
  import com.aizuda.snailjob.server.job.task.dto.JobTaskFailAlarmEventDTO;
  import com.aizuda.snailjob.server.job.task.dto.TaskExecuteDTO;
  import com.aizuda.snailjob.server.job.task.dto.WorkflowNodeTaskExecuteDTO;
  import com.aizuda.snailjob.server.job.task.support.JobExecutor;
  import com.aizuda.snailjob.server.job.task.support.JobTaskConverter;
  import com.aizuda.snailjob.server.job.task.support.alarm.event.JobTaskFailAlarmEvent;
  import com.aizuda.snailjob.server.job.task.support.executor.job.JobExecutorContext;
  import com.aizuda.snailjob.server.job.task.support.executor.job.JobExecutorFactory;
  import com.aizuda.snailjob.server.job.task.support.generator.task.JobTaskGenerateContext;
  import com.aizuda.snailjob.server.job.task.support.generator.task.JobTaskGenerator;
  import com.aizuda.snailjob.server.job.task.support.generator.task.JobTaskGeneratorFactory;
  import com.aizuda.snailjob.server.job.task.support.handler.JobTaskBatchHandler;
  import com.aizuda.snailjob.server.job.task.support.handler.WorkflowBatchHandler;
  import com.aizuda.snailjob.server.job.task.support.timer.JobTimeoutCheckTask;
  import com.aizuda.snailjob.server.job.task.support.timer.JobTimerTask;
  import com.aizuda.snailjob.server.job.task.support.timer.JobTimerWheel;
  import com.aizuda.snailjob.template.datasource.persistence.mapper.JobMapper;
  import com.aizuda.snailjob.template.datasource.persistence.mapper.JobTaskBatchMapper;
  import com.aizuda.snailjob.template.datasource.persistence.mapper.WorkflowTaskBatchMapper;
  import com.aizuda.snailjob.template.datasource.persistence.po.Job;
  import com.aizuda.snailjob.template.datasource.persistence.po.JobTask;
  import com.aizuda.snailjob.template.datasource.persistence.po.JobTaskBatch;
  import com.aizuda.snailjob.template.datasource.persistence.po.WorkflowTaskBatch;
  import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
  import com.google.common.collect.Lists;
  import lombok.RequiredArgsConstructor;
  import lombok.extern.slf4j.Slf4j;
  import org.jetbrains.annotations.NotNull;
  import org.springframework.beans.factory.config.ConfigurableBeanFactory;
  import org.springframework.context.annotation.Scope;
  import org.springframework.stereotype.Component;
  import org.springframework.transaction.TransactionStatus;
  import org.springframework.transaction.support.TransactionCallbackWithoutResult;
  import org.springframework.transaction.support.TransactionSynchronization;
  import org.springframework.transaction.support.TransactionSynchronizationManager;
  import org.springframework.transaction.support.TransactionTemplate;
  
  import java.text.MessageFormat;
  import java.time.Duration;
  import java.time.LocalDateTime;
  import java.util.List;
  import java.util.Objects;
  import java.util.Set;
  
  import static com.aizuda.snailjob.common.core.enums.JobTaskTypeEnum.MAP;
  import static com.aizuda.snailjob.common.core.enums.JobTaskTypeEnum.MAP_REDUCE;
  
  /**
   * @author: opensnail
   * @date : 2023-09-25 17:41
   * @since : 2.4.0
   */
  @Component(ActorGenerator.JOB_EXECUTOR_ACTOR)
  @Scope(ConfigurableBeanFactory.SCOPE_PROTOTYPE)
  @Slf4j
  @RequiredArgsConstructor
  public class JobExecutorActor extends AbstractActor {
      private final JobMapper jobMapper;
      private final JobTaskBatchMapper jobTaskBatchMapper;
      private final TransactionTemplate transactionTemplate;
      private final WorkflowBatchHandler workflowBatchHandler;
      private final JobTaskBatchHandler jobTaskBatchHandler;
      private final WorkflowTaskBatchMapper workflowTaskBatchMapper;
      private final InstanceManager instanceManager;
  
      @Override
      public Receive createReceive() {
          return receiveBuilder().match(TaskExecuteDTO.class, taskExecute -> {
              try {
                  log.debug("Preparing to execute task. [{}] [{}]", LocalDateTime.now(), JsonUtil.toJsonString(taskExecute));
                  transactionTemplate.execute(new TransactionCallbackWithoutResult() {
                      @Override
                      protected void doInTransactionWithoutResult(final TransactionStatus status) {
                          doExecute(taskExecute);
                      }
                  });
  
              } catch (Exception e) {
                  SnailJobLog.LOCAL.error("job executor exception. [{}]", taskExecute, e);
                  handleTaskBatch(taskExecute, JobTaskBatchStatusEnum.FAIL.getStatus(), JobOperationReasonEnum.TASK_EXECUTION_ERROR.getReason());
              } finally {
                  getContext().stop(getSelf());
              }
          }).build();
      }
  
      private void doExecute(final TaskExecuteDTO taskExecute) {
  
          LambdaQueryWrapper<Job> queryWrapper = new LambdaQueryWrapper<>();
          // 自动地校验任务必须是开启状态，手动触发无需校验
          if (JobTaskExecutorSceneEnum.AUTO_JOB.getType().equals(taskExecute.getTaskExecutorScene())) {
              queryWrapper.eq(Job::getJobStatus, StatusEnum.YES.getStatus());
          }
  
          Job job = jobMapper.selectOne(queryWrapper.eq(Job::getId, taskExecute.getJobId()));
          int taskStatus = JobTaskBatchStatusEnum.RUNNING.getStatus();
          try {
              int operationReason = JobOperationReasonEnum.NONE.getReason();
              if (Objects.isNull(job)) {
                  taskStatus = JobTaskBatchStatusEnum.CANCEL.getStatus();
                  operationReason = JobOperationReasonEnum.JOB_CLOSED.getReason();
              } else {
                  Set<InstanceLiveInfo> liveInfoSet = instanceManager.getInstanceALiveInfoSet(
                          job.getNamespaceId(), job.getGroupName(), job.getLabels());
  
                  // 空节点处理
                  if (CollUtil.isEmpty(liveInfoSet)) {
  
                      taskStatus = JobTaskBatchStatusEnum.CANCEL.getStatus();
                      operationReason = JobOperationReasonEnum.NOT_CLIENT.getReason();
  
                      WorkflowNodeTaskExecuteDTO taskExecuteDTO = new WorkflowNodeTaskExecuteDTO();
                      taskExecuteDTO.setWorkflowTaskBatchId(taskExecute.getWorkflowTaskBatchId());
                      taskExecuteDTO.setTaskExecutorScene(taskExecute.getTaskExecutorScene());
                      taskExecuteDTO.setParentId(taskExecute.getWorkflowNodeId());
                      taskExecuteDTO.setTaskBatchId(taskExecute.getTaskBatchId());
                      workflowBatchHandler.openNextNode(taskExecuteDTO);
  
                      // 无客户端节点-告警通知
                      SnailSpringContext.getContext().publishEvent(
                              new JobTaskFailAlarmEvent(JobTaskFailAlarmEventDTO.builder()
                                      .jobTaskBatchId(taskExecute.getTaskBatchId())
                                      .reason(JobNotifySceneEnum.JOB_NO_CLIENT_NODES_ERROR.getDesc())
                                      .notifyScene(JobNotifySceneEnum.JOB_NO_CLIENT_NODES_ERROR.getNotifyScene())
                                      .build()));
                  }
              }
  
              // 更新状态
              handleTaskBatch(taskExecute, taskStatus, operationReason);
  
              // 不是运行中的，不需要生产任务
              if (taskStatus != JobTaskBatchStatusEnum.RUNNING.getStatus()) {
                  return;
              }
  
              // 生成任务
              JobTaskGenerator taskInstance = JobTaskGeneratorFactory.getTaskInstance(job.getTaskType());
              JobTaskGenerateContext instanceGenerateContext = JobTaskConverter.INSTANCE.toJobTaskInstanceGenerateContext(job);
              instanceGenerateContext.setTaskBatchId(taskExecute.getTaskBatchId());
              if (Objects.nonNull(taskExecute.getTmpArgsStr())) {
                  instanceGenerateContext.setArgsStr(taskExecute.getTmpArgsStr());
              }
              if (Lists.newArrayList(MAP_REDUCE.getType(), MAP.getType()).contains(job.getTaskType())) {
                  instanceGenerateContext.setTaskName(SystemConstants.ROOT_MAP);
                  instanceGenerateContext.setMapSubTask(Lists.newArrayList(StrUtil.EMPTY));
                  instanceGenerateContext.setMrStage(MapReduceStageEnum.MAP.getStage());
              }
              List<JobTask> taskList = taskInstance.generate(instanceGenerateContext);
              if (CollUtil.isEmpty(taskList)) {
                  SnailJobLog.LOCAL.warn("Generate job task is empty, taskBatchId:[{}]", taskExecute.getTaskBatchId());
                  return;
              }
  
              // 事务提交以后再执行任务
              TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                  @Override
                  public void afterCommit() {
                      // 获取工作流的上下文
                      WorkflowTaskBatch workflowTaskBatch = null;
                      Long workflowTaskBatchId = taskExecute.getWorkflowTaskBatchId();
                      if (Objects.nonNull(workflowTaskBatchId)) {
                          workflowTaskBatch = workflowTaskBatchMapper.selectOne(
                                  new LambdaQueryWrapper<WorkflowTaskBatch>()
                                          .select(WorkflowTaskBatch::getWfContext)
                                          .eq(WorkflowTaskBatch::getId, taskExecute.getWorkflowTaskBatchId())
                          );
                      }
  
                      // 执行任务
                      JobExecutor jobExecutor = JobExecutorFactory.getJobExecutor(job.getTaskType());
                      jobExecutor.execute(buildJobExecutorContext(taskExecute, job, taskList, workflowTaskBatch));
                  }
              });
  
          } finally {
              log.debug("Task preparation complete.[{}]", JsonUtil.toJsonString(taskExecute));
              final int finalTaskStatus = taskStatus;
              TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                  @Override
                  public void afterCompletion(int status) {
                      // 清除时间轮的缓存
                      JobTimerWheel.clearCache(MessageFormat.format(JobTimerTask.IDEMPOTENT_KEY_PREFIX, taskExecute.getTaskBatchId()));
  
                      if (JobTaskBatchStatusEnum.RUNNING.getStatus() == finalTaskStatus) {
  
                          // 运行中的任务，需要进行超时检查
                          JobTimerWheel.registerWithJob(() -> new JobTimeoutCheckTask(taskExecute.getTaskBatchId(), job.getId()),
                                  // 加500ms是为了让尽量保证客户端自己先超时中断，防止客户端上报成功但是服务端已触发超时中断
                                  Duration.ofMillis(DateUtils.toEpochMilli(job.getExecutorTimeout()) + 500));
                      }
  
                      // 开启下一个常驻任务
                      jobTaskBatchHandler.openResidentTask(job, taskExecute);
                  }
              });
          }
  
      }
  
      @NotNull
      private static JobExecutorContext buildJobExecutorContext(TaskExecuteDTO taskExecute, Job job, List<JobTask> taskList,
                                                                final WorkflowTaskBatch workflowTaskBatch) {
          JobExecutorContext context = JobTaskConverter.INSTANCE.toJobExecutorContext(job);
          context.setTaskList(taskList);
          context.setTaskBatchId(taskExecute.getTaskBatchId());
          context.setJobId(job.getId());
          context.setWorkflowTaskBatchId(taskExecute.getWorkflowTaskBatchId());
          context.setWorkflowNodeId(taskExecute.getWorkflowNodeId());
          if (Objects.nonNull(workflowTaskBatch)) {
              context.setWfContext(workflowTaskBatch.getWfContext());
          }
          return context;
      }
  
      private void handleTaskBatch(TaskExecuteDTO taskExecute, int taskStatus, int operationReason) {
  
          JobTaskBatch jobTaskBatch = new JobTaskBatch();
          jobTaskBatch.setId(taskExecute.getTaskBatchId());
          jobTaskBatch.setExecutionAt(DateUtils.toNowMilli());
          jobTaskBatch.setTaskBatchStatus(taskStatus);
          jobTaskBatch.setOperationReason(operationReason);
          Assert.isTrue(1 == jobTaskBatchMapper.updateById(jobTaskBatch),
                  () -> new SnailJobServerException("Updating task failed"));
  
          if (JobTaskBatchStatusEnum.NOT_SUCCESS.contains(taskStatus)) {
              SnailSpringContext.getContext().publishEvent(
                      new JobTaskFailAlarmEvent(JobTaskFailAlarmEventDTO.builder()
                              .jobTaskBatchId(taskExecute.getTaskBatchId())
                              .reason(JobOperationReasonEnum.TASK_EXECUTION_ERROR.getDesc())
                              .notifyScene(JobNotifySceneEnum.JOB_TASK_ERROR.getNotifyScene())
                              .build()));
          }
      }
  
  }
  