Merge pull request #399 from dromara/dev

[ISSUE #398] fix tid confusion during alarming
dromara · Feb 8, 2024 · 515994f · 515994f
2 parents 7ac3a74 + 12b672a
commit 515994f
Showing 1 changed file with 20 additions and 6 deletions.
diff --git a/core/src/main/java/org/dromara/dynamictp/core/notifier/manager/AlarmManager.java b/core/src/main/java/org/dromara/dynamictp/core/notifier/manager/AlarmManager.java
@@ -78,11 +78,13 @@ public static void initAlarm(String poolName, NotifyItem notifyItem) {
         AlarmCounter.init(poolName, notifyItem.getType());
     }
 
-    public static void tryAlarmAsync(ExecutorWrapper executorWrapper, NotifyItemEnum notifyType, Runnable currRunnable) {
-        if (currRunnable instanceof DtpRunnable) {
-            MDC.put(TRACE_ID, ((DtpRunnable) currRunnable).getTraceId());
+    public static void tryAlarmAsync(ExecutorWrapper executorWrapper, NotifyItemEnum notifyType, Runnable runnable) {
+        preAlarm(runnable);
+        try {
+            ALARM_EXECUTOR.execute(() -> doTryAlarm(executorWrapper, notifyType));
+        } finally {
+            postAlarm(runnable);
         }
-        ALARM_EXECUTOR.execute(() -> doTryAlarm(executorWrapper, notifyType));
     }
 
     public static void tryAlarmAsync(ExecutorWrapper executorWrapper, List<NotifyItemEnum> notifyTypes) {
@@ -114,6 +116,10 @@ public static boolean checkThreshold(ExecutorWrapper executor, NotifyItemEnum no
         }
     }
 
+    public static void destroy() {
+        ALARM_EXECUTOR.shutdownNow();
+    }
+
     private static boolean checkLiveness(ExecutorWrapper executorWrapper, NotifyItem notifyItem) {
         val executor = executorWrapper.getExecutor();
         int maximumPoolSize = executor.getMaximumPoolSize();
@@ -136,7 +142,15 @@ private static boolean checkWithAlarmInfo(ExecutorWrapper executorWrapper, Notif
         return alarmInfo.getCount() >= notifyItem.getThreshold();
     }
 
-    public static void destroy() {
-        ALARM_EXECUTOR.shutdownNow();
+    private static void preAlarm(Runnable runnable) {
+        if (runnable instanceof DtpRunnable) {
+            MDC.put(TRACE_ID, ((DtpRunnable) runnable).getTraceId());
+        }
+    }
+
+    private static void postAlarm(Runnable runnable) {
+        if (runnable instanceof DtpRunnable) {
+            MDC.remove(TRACE_ID);
+        }
     }
 }