perf: 分库分表改造-任务表以及子表加入 task_instance_id 字段 #2991

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/config/JobFunctionConfiguration.java

@@ -24,7 +24,6 @@
 
 package com.tencent.bk.job.execute.config;
 
-import com.tencent.bk.job.common.util.json.JsonUtils;
 import com.tencent.bk.job.execute.engine.listener.CallbackListener;
 import com.tencent.bk.job.execute.engine.listener.GseTaskListener;
 import com.tencent.bk.job.execute.engine.listener.JobListener;
@@ -36,7 +35,6 @@
 import com.tencent.bk.job.execute.engine.listener.event.ResultHandleTaskResumeEvent;
 import com.tencent.bk.job.execute.engine.listener.event.StepEvent;
 import com.tencent.bk.job.execute.engine.model.JobCallbackDTO;
-import com.tencent.bk.job.execute.model.TaskInstanceRecordStateDO;
 import com.tencent.bk.job.execute.model.TaskNotifyDTO;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -58,31 +56,26 @@ public class JobFunctionConfiguration {
     public Consumer<Message<JobEvent>> handleJobEvent(@Autowired JobListener jobListener) {
         log.info("Init handleJobEvent consumer");
 
-        return message -> {
-            TaskInstanceRecordStateDO taskInstanceRecordStateDO = JsonUtils.fromJson(
-                message.getHeaders().get("TaskInstanceRecordStateDO", String.class),
-                TaskInstanceRecordStateDO.class);
-            jobListener.handleEvent(message.getPayload(), taskInstanceRecordStateDO);
-        };
+        return jobListener::onEvent;
     }
 
     @Bean
-    public Consumer<StepEvent> handleStepEvent(@Autowired StepListener stepListener) {
+    public Consumer<Message<StepEvent>> handleStepEvent(@Autowired StepListener stepListener) {
         log.info("Init handleStepEvent consumer");
-        return stepListener::handleEvent;
+        return stepListener::onEvent;
     }
 
     @Bean
-    public Consumer<GseTaskEvent> handleGseTaskEvent(@Autowired GseTaskListener gseTaskListener) {
+    public Consumer<Message<GseTaskEvent>> handleGseTaskEvent(@Autowired GseTaskListener gseTaskListener) {
         log.info("Init handleGseTaskEvent consumer");
-        return gseTaskListener::handleEvent;
+        return gseTaskListener::onEvent;
     }
 
     @Bean
-    public Consumer<ResultHandleTaskResumeEvent> handleResultHandleResumeEvent(
+    public Consumer<Message<ResultHandleTaskResumeEvent>> handleResultHandleResumeEvent(
         @Autowired ResultHandleResumeListener resultHandleResumeListener) {
         log.info("Init handleResultHandleResumeEvent consumer");
-        return resultHandleResumeListener::handleEvent;
+        return resultHandleResumeListener::onEvent;
     }
 
     @Bean

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/engine/listener/BaseJobExecuteMqListener.java

@@ -0,0 +1,66 @@
+/*
+ * Tencent is pleased to support the open source community by making BK-JOB蓝鲸智云作业平台 available.
+ *
+ * Copyright (C) 2021 THL A29 Limited, a Tencent company.  All rights reserved.
+ *
+ * BK-JOB蓝鲸智云作业平台 is licensed under the MIT License.
+ *
+ * License for BK-JOB蓝鲸智云作业平台:
+ * --------------------------------------------------------------------
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+ * documentation files (the "Software"), to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
+ * to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all copies or substantial portions of
+ * the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
+ * THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
+ * CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ */
+
+package com.tencent.bk.job.execute.engine.listener;
+
+import com.tencent.bk.job.common.util.json.JsonUtils;
+import com.tencent.bk.job.execute.colddata.JobExecuteContextThreadLocalRepo;
+import com.tencent.bk.job.execute.common.context.PropagatedJobExecuteContext;
+import com.tencent.bk.job.execute.engine.listener.event.Event;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+
+@Slf4j
+public abstract class BaseJobExecuteMqListener {
+
+
+    public final void onEvent(Message<? extends Event> message) {
+        beforeHandleMessage(message);
+        try {
+            handleEvent(message);
+        } finally {
+            afterHandle(message);
+        }
+    }
+
+    private void beforeHandleMessage(Message<?> message) {
+        log.info("beforeHandleMessage");
+        MessageHeaders headers = message.getHeaders();
+        String jobExecuteContextJson = (String) headers.get(PropagatedJobExecuteContext.KEY);
+        if (StringUtils.isNotEmpty(jobExecuteContextJson)) {
+            log.info("setJobExecuteContextThreadLocalRepo");
+            JobExecuteContextThreadLocalRepo.set(JsonUtils.fromJson(jobExecuteContextJson,
+                PropagatedJobExecuteContext.class));
+        }
+    }
+
+    private void afterHandle(Message<?> message) {
+        log.info("afterHandleMessage");
+        JobExecuteContextThreadLocalRepo.unset();
+    }
+
+    protected abstract void handleEvent(Message<? extends Event> message);
+}

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/engine/listener/GseTaskListener.java

@@ -28,20 +28,22 @@
 import com.tencent.bk.job.execute.common.exception.MessageHandlerUnavailableException;
 import com.tencent.bk.job.execute.engine.consts.GseTaskActionEnum;
 import com.tencent.bk.job.execute.engine.executor.GseTaskManager;
+import com.tencent.bk.job.execute.engine.listener.event.Event;
 import com.tencent.bk.job.execute.engine.listener.event.GseTaskEvent;
 import com.tencent.bk.job.execute.model.GseTaskDTO;
 import com.tencent.bk.job.execute.monitor.metrics.GseTasksExceptionCounter;
 import com.tencent.bk.job.execute.service.GseTaskService;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.Message;
 import org.springframework.stereotype.Component;
 
 /**
  * 执行引擎事件处理-GSE任务
  */
 @Component
 @Slf4j
-public class GseTaskListener {
+public class GseTaskListener extends BaseJobExecuteMqListener {
     private final GseTaskManager gseTaskManager;
     private final GseTaskService gseTaskService;
     private final GseTasksExceptionCounter gseTasksExceptionCounter;
@@ -58,9 +60,11 @@ public GseTaskListener(GseTaskManager gseTaskManager,
     /**
      * 处理GSE任务相关的事件
      *
-     * @param gseTaskEvent GSE任务事件
+     * @param message 消息
      */
-    public void handleEvent(GseTaskEvent gseTaskEvent) {
+    @Override
+    public void handleEvent(Message<? extends Event> message) {
+        GseTaskEvent gseTaskEvent = (GseTaskEvent) message.getPayload();
         log.info("Handle gse task event: {}, duration: {}ms", gseTaskEvent, gseTaskEvent.duration());
         GseTaskDTO gseTask = gseTaskService.getGseTask(gseTaskEvent.getJobInstanceId(), gseTaskEvent.getGseTaskId());
         String requestId = gseTaskEvent.getRequestId();

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/engine/listener/JobListener.java

@@ -30,14 +30,14 @@
 import com.tencent.bk.job.execute.common.constants.RunStatusEnum;
 import com.tencent.bk.job.execute.common.util.TaskCostCalculator;
 import com.tencent.bk.job.execute.engine.consts.JobActionEnum;
+import com.tencent.bk.job.execute.engine.listener.event.Event;
 import com.tencent.bk.job.execute.engine.listener.event.JobEvent;
 import com.tencent.bk.job.execute.engine.listener.event.StepEvent;
 import com.tencent.bk.job.execute.engine.listener.event.TaskExecuteMQEventDispatcher;
 import com.tencent.bk.job.execute.engine.model.JobCallbackDTO;
 import com.tencent.bk.job.execute.model.RollingConfigDTO;
 import com.tencent.bk.job.execute.model.StepInstanceBaseDTO;
 import com.tencent.bk.job.execute.model.TaskInstanceDTO;
-import com.tencent.bk.job.execute.model.TaskInstanceRecordStateDO;
 import com.tencent.bk.job.execute.model.db.RollingConfigDetailDO;
 import com.tencent.bk.job.execute.service.NotifyService;
 import com.tencent.bk.job.execute.service.RollingConfigService;
@@ -47,6 +47,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.Message;
 import org.springframework.stereotype.Component;
 
 import java.util.Collection;
@@ -57,7 +58,7 @@
  */
 @Component("jobEventListener")
 @Slf4j
-public class JobListener {
+public class JobListener extends BaseJobExecuteMqListener {
 
     private final TaskExecuteMQEventDispatcher taskExecuteMQEventDispatcher;
     private final StatisticsService statisticsService;
@@ -85,11 +86,12 @@ public JobListener(TaskExecuteMQEventDispatcher taskExecuteMQEventDispatcher,
     /**
      * 处理作业执行相关的事件
      *
-     * @param jobEvent 作业执行相关的事件
+     * @param message 消息
      */
-    public void handleEvent(JobEvent jobEvent, TaskInstanceRecordStateDO taskInstanceRecordStateDO) {
-        log.info("Handle job event, event: {}, duration: {}ms, taskInstanceRecordStateDO: {}",
-            jobEvent, jobEvent.duration(), taskInstanceRecordStateDO);
+    @Override
+    public void handleEvent(Message<? extends Event> message) {
+        JobEvent jobEvent = (JobEvent) message.getPayload();
+        log.info("Handle job event, event: {}, duration: {}ms", jobEvent, jobEvent.duration());
         long jobInstanceId = jobEvent.getJobInstanceId();
         JobActionEnum action = JobActionEnum.valueOf(jobEvent.getAction());
         try {
@@ -319,7 +321,7 @@ private void startStep(StepInstanceBaseDTO stepInstance) {
                 stepInstance.getId(), RunStatusEnum.ROLLING_WAITING.getValue());
             taskExecuteMQEventDispatcher.dispatchStepEvent(
                 StepEvent.startStep(stepInstance.getTaskInstanceId(), stepInstance.getId(),
-                stepInstance.getBatch() + 1));
+                    stepInstance.getBatch() + 1));
         } else {
             stepInstanceService.updateStepStatus(stepInstance.getTaskInstanceId(),
                 stepInstance.getId(), RunStatusEnum.BLANK.getValue());

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/engine/listener/ResultHandleResumeListener.java

@@ -30,6 +30,7 @@
 import com.tencent.bk.job.execute.common.constants.RunStatusEnum;
 import com.tencent.bk.job.execute.config.FileDistributeConfig;
 import com.tencent.bk.job.execute.engine.evict.TaskEvictPolicyExecutor;
+import com.tencent.bk.job.execute.engine.listener.event.Event;
 import com.tencent.bk.job.execute.engine.listener.event.ResultHandleTaskResumeEvent;
 import com.tencent.bk.job.execute.engine.listener.event.TaskExecuteMQEventDispatcher;
 import com.tencent.bk.job.execute.engine.model.FileDest;
@@ -56,6 +57,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.Message;
 import org.springframework.stereotype.Component;
 
 import java.util.HashMap;
@@ -69,7 +71,7 @@
  */
 @Component
 @Slf4j
-public class ResultHandleResumeListener {
+public class ResultHandleResumeListener extends BaseJobExecuteMqListener {
     private final TaskInstanceService taskInstanceService;
 
     private final ResultHandleManager resultHandleManager;
@@ -132,7 +134,8 @@ public ResultHandleResumeListener(TaskInstanceService taskInstanceService,
     /**
      * 恢复被中断的作业结果处理任务
      */
-    public void handleEvent(ResultHandleTaskResumeEvent event) {
+    public void handleEvent(Message<? extends Event> message) {
+        ResultHandleTaskResumeEvent event = (ResultHandleTaskResumeEvent) message.getPayload();
         log.info("Receive gse task result handle task resume event: {}, duration: {}ms", event, event.duration());
         GseTaskDTO gseTask = gseTaskService.getGseTask(event.getJobInstanceId(), event.getGseTaskId());
         long stepInstanceId = gseTask.getStepInstanceId();

src/backend/job-execute/service-job-execute/src/main/java/com/tencent/bk/job/execute/engine/listener/StepListener.java

@@ -25,19 +25,21 @@
 package com.tencent.bk.job.execute.engine.listener;
 
 import com.tencent.bk.job.execute.common.constants.StepExecuteTypeEnum;
+import com.tencent.bk.job.execute.engine.listener.event.Event;
 import com.tencent.bk.job.execute.engine.listener.event.StepEvent;
 import com.tencent.bk.job.execute.model.StepInstanceDTO;
 import com.tencent.bk.job.execute.service.StepInstanceService;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.Message;
 import org.springframework.stereotype.Component;
 
 /**
  * 执行引擎事件处理-步骤
  */
 @Component
 @Slf4j
-public class StepListener {
+public class StepListener extends BaseJobExecuteMqListener {
     private final StepInstanceService stepInstanceService;
     private final GseStepEventHandler gseStepEventHandler;
     private final ConfirmStepEventHandler confirmStepEventHandler;
@@ -54,9 +56,11 @@ public StepListener(StepInstanceService stepInstanceService,
     /**
      * 处理步骤执行相关的事件
      *
-     * @param stepEvent 步骤执行相关的事件
+     * @param message 消息
      */
-    public void handleEvent(StepEvent stepEvent) {
+    @Override
+    public void handleEvent(Message<? extends Event> message) {
+        StepEvent stepEvent = (StepEvent) message.getPayload();
         log.info("Handle step event: {}, duration: {}ms", stepEvent, stepEvent.duration());
         long stepInstanceId = stepEvent.getStepInstanceId();
         try {