From c2da994e31bec043b2a210a4e58f7785c8ca8166 Mon Sep 17 00:00:00 2001
From: "chen.ma" <machencoding@163.com>
Date: Thu, 10 Mar 2022 08:43:38 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B6=88=E6=81=AF=E6=8A=A5=E8=AD=A6=E6=8E=A8?=
 =?UTF-8?q?=E9=80=81=E4=BC=98=E5=8C=96.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../notify/BaseSendMessageServiceImpl.java    | 12 +++---
 .../platform/DingSendMessageHandler.java      |  4 +-
 .../platform/LarkSendMessageHandler.java      |  4 +-
 .../platform/WeChatSendMessageHandler.java    |  4 +-
 .../cn/hippo4j/common/toolkit/MDCUtil.java    | 25 ------------
 .../ThreadPoolNotifyAlarmHandler.java         | 16 +++-----
 .../core/toolkit/TraceContextUtil.java        | 40 +++++++++++++++++++
 7 files changed, 59 insertions(+), 46 deletions(-)
 delete mode 100644 hippo4j-common/src/main/java/cn/hippo4j/common/toolkit/MDCUtil.java
 create mode 100644 hippo4j-core/src/main/java/cn/hippo4j/core/toolkit/TraceContextUtil.java

diff --git a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/BaseSendMessageServiceImpl.java b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/BaseSendMessageServiceImpl.java
index eafe706f..7e26cff5 100644
--- a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/BaseSendMessageServiceImpl.java
+++ b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/BaseSendMessageServiceImpl.java
@@ -50,7 +50,8 @@ public class BaseSendMessageServiceImpl implements HippoSendMessageService, Comm
                     return;
                 }
 
-                if (isSendAlarm(each.getThreadPoolId(), each.setTypeEnum(typeEnum))) {
+                if (isSendAlarm(each.getThreadPoolId(), each.getPlatform(), typeEnum)) {
+                    alarmNotifyRequest.setNotifyTypeEnum(typeEnum);
                     messageHandler.sendAlarmMessage(each, alarmNotifyRequest);
                 }
             } catch (Exception ex) {
@@ -88,14 +89,15 @@ public class BaseSendMessageServiceImpl implements HippoSendMessageService, Comm
      * Is send alarm.
      *
      * @param threadPoolId
-     * @param notifyInfo
+     * @param platform
+     * @param typeEnum
      * @return
      */
-    private boolean isSendAlarm(String threadPoolId, NotifyConfigDTO notifyInfo) {
+    private boolean isSendAlarm(String threadPoolId, String platform, NotifyTypeEnum typeEnum) {
         AlarmControlDTO alarmControl = AlarmControlDTO.builder()
                 .threadPool(threadPoolId)
-                .platform(notifyInfo.getPlatform())
-                .typeEnum(notifyInfo.getTypeEnum())
+                .platform(platform)
+                .typeEnum(typeEnum)
                 .build();
 
         return alarmControlHandler.isSendAlarm(alarmControl);
diff --git a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/DingSendMessageHandler.java b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/DingSendMessageHandler.java
index 9feeffd4..059e782c 100644
--- a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/DingSendMessageHandler.java
+++ b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/DingSendMessageHandler.java
@@ -40,7 +40,7 @@ public class DingSendMessageHandler implements SendMessageHandler<AlarmNotifyReq
 
         String dingAlarmTxt;
         String dingAlarmTimoutReplaceTxt;
-        if (Objects.equals(notifyConfig.getTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
+        if (Objects.equals(alarmNotifyRequest.getNotifyTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
             String executeTimeoutTrace = alarmNotifyRequest.getExecuteTimeoutTrace();
             if (StringUtil.isNotBlank(executeTimeoutTrace)) {
                 String dingAlarmTimoutTraceReplaceTxt = String.format(DING_ALARM_TIMOUT_TRACE_REPLACE_TXT, executeTimeoutTrace);
@@ -61,7 +61,7 @@ public class DingSendMessageHandler implements SendMessageHandler<AlarmNotifyReq
                 // 环境
                 alarmNotifyRequest.getActive(),
                 // 报警类型
-                notifyConfig.getTypeEnum(),
+                alarmNotifyRequest.getNotifyTypeEnum(),
                 // 线程池ID
                 alarmNotifyRequest.getThreadPoolId(),
                 // 应用名称
diff --git a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/LarkSendMessageHandler.java b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/LarkSendMessageHandler.java
index 35bdf75b..0f63feb3 100644
--- a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/LarkSendMessageHandler.java
+++ b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/LarkSendMessageHandler.java
@@ -42,7 +42,7 @@ public class LarkSendMessageHandler implements SendMessageHandler<AlarmNotifyReq
 
         String larkAlarmTxt;
         String larkAlarmTimoutReplaceTxt;
-        if (Objects.equals(notifyConfig.getTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
+        if (Objects.equals(alarmNotifyRequest.getNotifyTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
             String executeTimeoutTrace = alarmNotifyRequest.getExecuteTimeoutTrace();
             if (StringUtil.isNotBlank(executeTimeoutTrace)) {
                 String larkAlarmTimoutTraceReplaceTxt = String.format(LARK_ALARM_TIMOUT_TRACE_REPLACE_TXT, executeTimeoutTrace);
@@ -61,7 +61,7 @@ public class LarkSendMessageHandler implements SendMessageHandler<AlarmNotifyReq
                 // 环境
                 alarmNotifyRequest.getActive(),
                 // 报警类型
-                notifyConfig.getTypeEnum(),
+                alarmNotifyRequest.getNotifyTypeEnum(),
                 // 线程池ID
                 alarmNotifyRequest.getThreadPoolId(),
                 // 应用名称
diff --git a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/WeChatSendMessageHandler.java b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/WeChatSendMessageHandler.java
index ee7fd89e..e912176a 100644
--- a/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/WeChatSendMessageHandler.java
+++ b/hippo4j-common/src/main/java/cn/hippo4j/common/notify/platform/WeChatSendMessageHandler.java
@@ -38,7 +38,7 @@ public class WeChatSendMessageHandler implements SendMessageHandler<AlarmNotifyR
 
         String weChatAlarmTxt;
         String weChatAlarmTimoutReplaceTxt;
-        if (Objects.equals(notifyConfig.getTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
+        if (Objects.equals(alarmNotifyRequest.getNotifyTypeEnum(), NotifyTypeEnum.TIMEOUT)) {
             String executeTimeoutTrace = alarmNotifyRequest.getExecuteTimeoutTrace();
             if (StringUtil.isNotBlank(executeTimeoutTrace)) {
                 String weChatAlarmTimoutTraceReplaceTxt = String.format(WE_CHAT_ALARM_TIMOUT_TRACE_REPLACE_TXT, executeTimeoutTrace);
@@ -58,7 +58,7 @@ public class WeChatSendMessageHandler implements SendMessageHandler<AlarmNotifyR
                 // 环境
                 alarmNotifyRequest.getActive(),
                 // 报警类型
-                notifyConfig.getTypeEnum(),
+                alarmNotifyRequest.getNotifyTypeEnum(),
                 // 线程池ID
                 alarmNotifyRequest.getThreadPoolId(),
                 // 应用名称
diff --git a/hippo4j-common/src/main/java/cn/hippo4j/common/toolkit/MDCUtil.java b/hippo4j-common/src/main/java/cn/hippo4j/common/toolkit/MDCUtil.java
deleted file mode 100644
index f72ec60d..00000000
--- a/hippo4j-common/src/main/java/cn/hippo4j/common/toolkit/MDCUtil.java
+++ /dev/null
@@ -1,25 +0,0 @@
-package cn.hippo4j.common.toolkit;
-
-import org.slf4j.MDC;
-
-/**
- * MD util.
- *
- * @author chen.ma
- * @date 2022/3/3 08:30
- */
-public class MDCUtil {
-
-    /**
-     * Get and remove.
-     *
-     * @param key
-     * @return
-     */
-    public static String getAndRemove(String key) {
-        String val = MDC.get(key);
-        MDC.remove(key);
-        return val;
-    }
-
-}
diff --git a/hippo4j-core/src/main/java/cn/hippo4j/core/executor/ThreadPoolNotifyAlarmHandler.java b/hippo4j-core/src/main/java/cn/hippo4j/core/executor/ThreadPoolNotifyAlarmHandler.java
index d6d04d88..62306bd9 100644
--- a/hippo4j-core/src/main/java/cn/hippo4j/core/executor/ThreadPoolNotifyAlarmHandler.java
+++ b/hippo4j-core/src/main/java/cn/hippo4j/core/executor/ThreadPoolNotifyAlarmHandler.java
@@ -5,7 +5,7 @@ import cn.hippo4j.common.notify.NotifyTypeEnum;
 import cn.hippo4j.common.notify.ThreadPoolNotifyAlarm;
 import cn.hippo4j.common.notify.request.AlarmNotifyRequest;
 import cn.hippo4j.common.notify.request.ChangeParameterNotifyRequest;
-import cn.hippo4j.common.toolkit.MDCUtil;
+import cn.hippo4j.core.toolkit.TraceContextUtil;
 import cn.hippo4j.common.toolkit.StringUtil;
 import cn.hippo4j.core.executor.manage.GlobalNotifyAlarmManage;
 import cn.hippo4j.core.executor.manage.GlobalThreadPoolManage;
@@ -22,8 +22,6 @@ import org.springframework.boot.CommandLineRunner;
 import java.util.List;
 import java.util.concurrent.*;
 
-import static cn.hippo4j.common.constant.Constants.EXECUTE_TIMEOUT_TRACE;
-
 /**
  * Thread pool alarm notify.
  *
@@ -168,15 +166,13 @@ public class ThreadPoolNotifyAlarmHandler implements Runnable, CommandLineRunner
                 alarmNotifyRequest.setThreadPoolId(threadPoolId);
                 alarmNotifyRequest.setExecuteTime(executeTime);
                 alarmNotifyRequest.setExecuteTimeOut(executeTimeOut);
-                String executeTimeoutTrace = MDCUtil.getAndRemove(EXECUTE_TIMEOUT_TRACE);
-                Runnable task = () -> {
-                    if (StringUtil.isNotBlank(executeTimeoutTrace)) {
-                        alarmNotifyRequest.setExecuteTimeoutTrace(executeTimeoutTrace);
-                    }
 
-                    hippoSendMessageService.sendAlarmMessage(NotifyTypeEnum.TIMEOUT, alarmNotifyRequest);
-                };
+                String executeTimeoutTrace = TraceContextUtil.getAndRemove();
+                if (StringUtil.isNotBlank(executeTimeoutTrace)) {
+                    alarmNotifyRequest.setExecuteTimeoutTrace(executeTimeoutTrace);
+                }
 
+                Runnable task = () -> hippoSendMessageService.sendAlarmMessage(NotifyTypeEnum.TIMEOUT, alarmNotifyRequest);
                 EXECUTE_TIMEOUT_EXECUTOR.execute(task);
             } catch (Throwable ex) {
                 log.error("Send thread pool execution timeout alarm error.", ex);
diff --git a/hippo4j-core/src/main/java/cn/hippo4j/core/toolkit/TraceContextUtil.java b/hippo4j-core/src/main/java/cn/hippo4j/core/toolkit/TraceContextUtil.java
new file mode 100644
index 00000000..101865fe
--- /dev/null
+++ b/hippo4j-core/src/main/java/cn/hippo4j/core/toolkit/TraceContextUtil.java
@@ -0,0 +1,40 @@
+package cn.hippo4j.core.toolkit;
+
+import org.slf4j.MDC;
+
+import static cn.hippo4j.common.constant.Constants.EXECUTE_TIMEOUT_TRACE;
+
+/**
+ * MD util.
+ *
+ * @author chen.ma
+ * @date 2022/3/3 08:30
+ */
+public class TraceContextUtil {
+
+    /**
+     * Execute timeout trace key.
+     */
+    private static String EXECUTE_TIMEOUT_TRACE_KEY = EXECUTE_TIMEOUT_TRACE;
+
+    /**
+     * Get and remove.
+     *
+     * @return
+     */
+    public static String getAndRemove() {
+        String val = MDC.get(EXECUTE_TIMEOUT_TRACE_KEY);
+        MDC.remove(EXECUTE_TIMEOUT_TRACE_KEY);
+        return val;
+    }
+
+    /**
+     * Set execute timeout trace key.
+     *
+     * @param key
+     */
+    public static void setExecuteTimeoutTraceKey(String key) {
+        EXECUTE_TIMEOUT_TRACE_KEY = key;
+    }
+
+}