Merge branch 'master' of https://github.com/jovezhao/nest

jovezhao · Feb 9, 2022 · 3e47136 · 3e47136
2 parents 04784dc + 35a9263
commit 3e47136
Showing 1 changed file with 60 additions and 13 deletions.
diff --git a/nest-ddd/src/main/java/com/zhaofujun/nest/context/appservice/ServiceContextManager.java b/nest-ddd/src/main/java/com/zhaofujun/nest/context/appservice/ServiceContextManager.java
@@ -3,30 +3,77 @@
 import com.alibaba.ttl.TransmittableThreadLocal;
 
 import java.util.Stack;
+import java.util.concurrent.ConcurrentHashMap;
 
 public class ServiceContextManager {
 
-    //使用堆栈来处理当前上下文信息
+    /**
+     * 使用堆栈来处理当前上下文信息
+     */
+    private static TransmittableThreadLocal<ConcurrentHashMap<String, Stack<ServiceContext>>> serviceContextThreadLocal = new TransmittableThreadLocal<>();
+
+    /**
+     * 用来记录主线程名称,主要用于应用服务开启多线程调用领域服务方法，通过记录的主线程名来获取主线程的上下文，用于提交领域服务的事务
+     */
+    private static TransmittableThreadLocal<String> mainNameThreadLocal = new TransmittableThreadLocal<>();
 
-    private static TransmittableThreadLocal<Stack<ServiceContext>> serviceContextThreadLocal = new TransmittableThreadLocal<>();
 
     public static void set(ServiceContext serviceContext) {
-        Stack<ServiceContext> serviceContexts = serviceContextThreadLocal.get();
-        if (serviceContexts == null) {
-            serviceContexts = new Stack<>();
-            serviceContextThreadLocal.set(serviceContexts);
+        ConcurrentHashMap<String, Stack<ServiceContext>> serviceContextHashMap = serviceContextThreadLocal.get();
+        if (serviceContextHashMap == null) {
+            serviceContextHashMap = new ConcurrentHashMap<>();
+            serviceContextThreadLocal.set(serviceContextHashMap);
+        }
+        if (serviceContextHashMap.isEmpty()) {
+            mainNameThreadLocal.set(Thread.currentThread().getName());
         }
-        serviceContexts.push(serviceContext);
+        Stack<ServiceContext> serviceContextStack = serviceContextHashMap.get(Thread.currentThread().getName());
+        if (serviceContextStack == null) {
+            serviceContextStack = new Stack<>();
+        }
+        serviceContextStack.push(serviceContext);
+        serviceContextHashMap.put(Thread.currentThread().getName(), serviceContextStack);
     }
 
     public static ServiceContext get() {
-        Stack<ServiceContext> serviceContexts = serviceContextThreadLocal.get();
-        if (serviceContexts == null || serviceContexts.empty()) return null;
-        return serviceContexts.peek();
+        Stack<ServiceContext> serviceContextStack = getStack();
+        if (serviceContextStack == null || serviceContextStack.empty()) {
+            return null;
+        }
+        return serviceContextStack.peek();
+    }
+
+    public static Stack<ServiceContext> getStack() {
+        ConcurrentHashMap<String, Stack<ServiceContext>> serviceContextHashMap = serviceContextThreadLocal.get();
+        if (serviceContextHashMap == null || serviceContextHashMap.isEmpty()) {
+            return null;
+        }
+        Stack<ServiceContext> serviceContextStack = serviceContextHashMap.get(Thread.currentThread().getName());
+        if (serviceContextStack == null || serviceContextStack.empty()) {
+            serviceContextStack = serviceContextHashMap.get(mainNameThreadLocal.get());
+        }
+        return serviceContextStack;
     }
 
-    public static ServiceContext pop() {
-        Stack<ServiceContext> serviceContexts = serviceContextThreadLocal.get();
-        return serviceContexts.pop();
+    public static void pop() {
+        ConcurrentHashMap<String, Stack<ServiceContext>> serviceContextHashMap = serviceContextThreadLocal.get();
+        if (serviceContextHashMap == null || serviceContextHashMap.isEmpty()) {
+            serviceContextThreadLocal.remove();
+            mainNameThreadLocal.remove();
+            return;
+        }
+        Stack<ServiceContext> serviceContextStack = serviceContextHashMap.get(Thread.currentThread().getName());
+        if (serviceContextStack != null && !serviceContextStack.isEmpty()) {
+            serviceContextStack.pop();
+        }
+        if (serviceContextStack == null || serviceContextStack.isEmpty()) {
+            serviceContextHashMap.remove(Thread.currentThread().getName());
+        }
+        if (serviceContextHashMap.isEmpty()) {
+            serviceContextThreadLocal.remove();
+            mainNameThreadLocal.remove();
+        }
     }
+
+
 }