release threadLocal store

oh-my-scheduler-worker/src/main/java/com/github/kfcfans/oms/worker/common/ThreadLocalStore.java

@@ -6,15 +6,36 @@ import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * 存储一些不方便直接传递的东西
- * #attention：警惕内存泄漏问题，最好在 ProcessorTracker destroy 时，执行 remove
+ * #attention：警惕内存泄漏问题，执行完毕后手动释放
  *
  * @author tjq
  * @since 2020/3/18
  */
 public class ThreadLocalStore {
 
-    public static final ThreadLocal<TaskDO> TASK_THREAD_LOCAL = new ThreadLocal<>();
+    private static final ThreadLocal<TaskDO> TASK_THREAD_LOCAL = new ThreadLocal<>();
 
-    public static final ThreadLocal<AtomicLong> TASK_ID_THREAD_LOCAL = new ThreadLocal<>();
+    private static final ThreadLocal<AtomicLong> TASK_ID_THREAD_LOCAL = new ThreadLocal<>();
+
+
+    public static TaskDO getTask() {
+        return TASK_THREAD_LOCAL.get();
+    }
+
+    public static void setTask(TaskDO task) {
+        TASK_THREAD_LOCAL.set(task);
+    }
+
+    public static AtomicLong getTaskIDAddr() {
+        if (TASK_ID_THREAD_LOCAL.get() == null) {
+            TASK_ID_THREAD_LOCAL.set(new AtomicLong(0));
+        }
+        return TASK_ID_THREAD_LOCAL.get();
+    }
+
+    public static void clear() {
+        TASK_ID_THREAD_LOCAL.remove();
+        TASK_THREAD_LOCAL.remove();
+    }
 
 }

oh-my-scheduler-worker/src/main/java/com/github/kfcfans/oms/worker/core/executor/ProcessorRunnable.java

@@ -21,12 +21,10 @@ import com.github.kfcfans.oms.worker.sdk.api.BroadcastProcessor;
 import com.github.kfcfans.oms.worker.sdk.api.MapReduceProcessor;
 import com.google.common.base.Stopwatch;
 import lombok.AllArgsConstructor;
-import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.BeanUtils;
 
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Processor 执行器
@@ -39,20 +37,17 @@ import java.util.concurrent.atomic.AtomicLong;
 public class ProcessorRunnable implements Runnable {
 
 
-    private InstanceInfo instanceInfo;
+    private final InstanceInfo instanceInfo;
     private final ActorSelection taskTrackerActor;
-    @Getter
     private final TaskDO task;
 
-    @Override
-    public void run() {
+    public void innerRun() {
 
         String taskId = task.getTaskId();
         String instanceId = task.getInstanceId();
 
         log.debug("[ProcessorRunnable-{}] start to run task(taskId={}&taskName={})", instanceId, taskId, task.getTaskName());
 
-        try {
         // 0. 完成执行上下文准备 & 上报执行信息
         TaskContext taskContext = new TaskContext();
         BeanUtils.copyProperties(task, taskContext);
@@ -63,8 +58,7 @@ public class ProcessorRunnable implements Runnable {
         if (task.getTaskContent() != null && task.getTaskContent().length > 0) {
             taskContext.setSubTask(SerializerUtils.deSerialized(task.getTaskContent()));
         }
-            ThreadLocalStore.TASK_THREAD_LOCAL.set(task);
-            ThreadLocalStore.TASK_ID_THREAD_LOCAL.set(new AtomicLong(0));
+        ThreadLocalStore.setTask(task);
 
         reportStatus(TaskStatus.WORKER_PROCESSING, null);
 
@@ -150,10 +144,6 @@ public class ProcessorRunnable implements Runnable {
             processResult = new ProcessResult(false, e.toString());
         }
         reportStatus(processResult.isSuccess() ? TaskStatus.WORKER_PROCESS_SUCCESS : TaskStatus.WORKER_PROCESS_FAILED, processResult.getMsg());
-
-        }catch (Exception e) {
-            log.error("[ProcessorRunnable-{}] execute failed, please fix this bug!", instanceId, e);
-        }
     }
 
     private BasicProcessor getProcessor() {
@@ -193,4 +183,15 @@ public class ProcessorRunnable implements Runnable {
 
         taskTrackerActor.tell(req, null);
     }
+
+    @Override
+    public void run() {
+        try {
+            innerRun();
+        }catch (Exception e) {
+            log.error("[ProcessorRunnable-{}] execute failed, please fix this bug @tjq!", task.getInstanceId(), e);
+        }finally {
+            ThreadLocalStore.clear();
+        }
+    }
 }

oh-my-scheduler-worker/src/main/java/com/github/kfcfans/oms/worker/pojo/request/ProcessorMapTaskRequest.java

@@ -42,7 +42,7 @@ public class ProcessorMapTaskRequest implements Serializable {
 
         subTaskList.forEach(subTask -> {
             // 同一个 Task 内部可能多次 Map，因此还是要确保线程级别的唯一
-            String subTaskId = parentTask.getTaskId() + "." + ThreadLocalStore.TASK_ID_THREAD_LOCAL.get().getAndIncrement();
+            String subTaskId = parentTask.getTaskId() + "." + ThreadLocalStore.getTaskIDAddr().getAndIncrement();
             // 写入类名，方便反序列化
             subTasks.add(new SubTask(subTaskId, SerializerUtils.serialize(subTask)));
         });

oh-my-scheduler-worker/src/main/java/com/github/kfcfans/oms/worker/sdk/api/MapReduceProcessor.java

@@ -49,7 +49,7 @@ public abstract class MapReduceProcessor implements BasicProcessor {
             log.warn("[MapReduceProcessor] map task size is too large, network maybe overload... please try to split the tasks.");
         }
 
-        TaskDO task = ThreadLocalStore.TASK_THREAD_LOCAL.get();
+        TaskDO task = ThreadLocalStore.getTask();
 
         // 1. 构造请求
         ProcessorMapTaskRequest req = new ProcessorMapTaskRequest(task, taskList, taskName);
@@ -74,7 +74,7 @@ public abstract class MapReduceProcessor implements BasicProcessor {
     }
 
     public boolean isRootTask() {
-        TaskDO task = ThreadLocalStore.TASK_THREAD_LOCAL.get();
+        TaskDO task = ThreadLocalStore.getTask();
         return TaskConstant.ROOT_TASK_ID.equals(task.getTaskId());
     }