Merge pull request Nepxion#3 from Nepxion/zifeihan

Avoid calling in the same thread
LittleCadet · Apr 27, 2021 · 7381fd1 · 7381fd1
2 parents 9bd50e8 + 512d4eb
commit 7381fd1
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/async/AsyncContext.java b/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/async/AsyncContext.java
@@ -12,11 +12,21 @@
 public class AsyncContext {
     private Object[] objects;
 
+    /**
+     * Record the original thread and compare it with the thread that called runnable/callable.
+     */
+    private Thread originThread;
+
     public AsyncContext(Object[] objects) {
         this.objects = objects;
+        this.originThread = Thread.currentThread();
     }
 
     public Object[] getObjects() {
         return objects;
     }
+
+    public Thread getOriginThread() {
+        return originThread;
+    }
 }
diff --git a/...t-starter/src/main/java/com/nepxion/discovery/agent/plugin/spring/async/WrapCallable.java b/...t-starter/src/main/java/com/nepxion/discovery/agent/plugin/spring/async/WrapCallable.java
@@ -27,6 +27,13 @@ public WrapCallable(Callable<T> callable) {
 
     @Override
     public T call() throws Exception {
+
+        //Avoid calling in the same thread.
+        if (asyncContext.getOriginThread().equals(Thread.currentThread())) {
+            return callable.call();
+        }
+
+        //Call the original method and copy some objects held by ThreadLocal.
         Object[] objects = asyncContext.getObjects();
         ThreadLocalCopier.before(objects);
         try {

diff --git a/...ain/java/com/nepxion/discovery/agent/plugin/thread/interceptor/ThreadCallInterceptor.java b/...ain/java/com/nepxion/discovery/agent/plugin/thread/interceptor/ThreadCallInterceptor.java
@@ -21,13 +21,26 @@ public static void before(Object object) {
             if (null == asyncContext) {
                 return;
             }
+            //Avoid calling in the same thread.
+            if (asyncContext.getOriginThread().equals(Thread.currentThread())) {
+                return;
+            }
             Object[] objects = asyncContext.getObjects();
             ThreadLocalCopier.before(objects);
         }
     }
 
     public static void after(Object object) {
         if (object instanceof AsyncContextAccessor) {
+            AsyncContextAccessor asyncContextAccessor = (AsyncContextAccessor) object;
+            AsyncContext asyncContext = asyncContextAccessor.getAsyncContext();
+            if (null == asyncContext) {
+                return;
+            }
+            //Avoid calling in the same thread.
+            if (asyncContext.getOriginThread().equals(Thread.currentThread())) {
+                return;
+            }
             ThreadLocalCopier.after();
         }
     }