优化LocalStoreageLifeCycle逻辑

优化gray mock filter

Author: saleson

spring-cloud-gray-client-netflix/src/main/java/cn/springcloud/gray/client/netflix/hystrix/HystrixLocalStorageCycle.java

@@ -2,27 +2,45 @@
 
 import cn.springcloud.gray.request.LocalStorageLifeCycle;
 import com.netflix.hystrix.strategy.concurrency.HystrixRequestContext;
+import org.apache.commons.lang3.StringUtils;
 
 public class HystrixLocalStorageCycle implements LocalStorageLifeCycle {
 
 
-    private ThreadLocal<Boolean> hystrixRequestContextInitialized = new ThreadLocal<>();
+    private ThreadLocal<String> hystrixRequestContextInitializedMark = new ThreadLocal<>();
 
 
     @Override
     public void initContext() {
+        initContext("");
+    }
+
+    @Override
+    public void initContext(String mark) {
         if (!HystrixRequestContext.isCurrentThreadInitialized()) {
             HystrixRequestContext.initializeContext();
-            hystrixRequestContextInitialized.set(true);
+            hystrixRequestContextInitializedMark.set(mark);
         }
     }
 
     @Override
     public void closeContext() {
-        Boolean hystrixReqCxtInited = hystrixRequestContextInitialized.get();
-        if (hystrixReqCxtInited != null) {
-            hystrixRequestContextInitialized.remove();
-            if (hystrixReqCxtInited && HystrixRequestContext.isCurrentThreadInitialized()) {
+        String hystrixReqCxtInitedMark = hystrixRequestContextInitializedMark.get();
+        if (hystrixReqCxtInitedMark != null) {
+            hystrixRequestContextInitializedMark.remove();
+            if (HystrixRequestContext.isCurrentThreadInitialized()) {
+                HystrixRequestContext.getContextForCurrentThread().shutdown();
+            }
+
+        }
+    }
+
+    @Override
+    public void closeContext(String mark) {
+        String hystrixReqCxtInitedMark = hystrixRequestContextInitializedMark.get();
+        if (hystrixReqCxtInitedMark != null && StringUtils.equals(mark, hystrixReqCxtInitedMark)) {
+            hystrixRequestContextInitializedMark.remove();
+            if (HystrixRequestContext.isCurrentThreadInitialized()) {
                 HystrixRequestContext.getContextForCurrentThread().shutdown();
             }
 

spring-cloud-gray-client-netflix/src/main/java/cn/springcloud/gray/client/netflix/hystrix/HystrixRequestLocalStorage.java

@@ -1,17 +1,27 @@
 package cn.springcloud.gray.client.netflix.hystrix;
 
+import cn.springcloud.gray.request.BaseRequestLocalStorage;
 import cn.springcloud.gray.request.GrayRequest;
 import cn.springcloud.gray.request.GrayTrackInfo;
-import cn.springcloud.gray.request.RequestLocalStorage;
+import cn.springcloud.gray.request.LocalStorageLifeCycle;
 import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariableDefault;
 
-public class HystrixRequestLocalStorage implements RequestLocalStorage {
+public class HystrixRequestLocalStorage extends BaseRequestLocalStorage {
 
 
     private static final HystrixRequestVariableDefault<GrayTrackInfo> grayTrackInfoLocal = new HystrixRequestVariableDefault<GrayTrackInfo>();
     private static final HystrixRequestVariableDefault<GrayRequest> grayRequestLocal = new HystrixRequestVariableDefault<GrayRequest>();
 
 
+    public HystrixRequestLocalStorage() {
+        this(new HystrixLocalStorageCycle());
+    }
+
+    public HystrixRequestLocalStorage(LocalStorageLifeCycle localStorageLifeCycle) {
+        super(localStorageLifeCycle);
+    }
+
+
     @Override
     public void setGrayTrackInfo(GrayTrackInfo grayTrackInfo) {
         grayTrackInfoLocal.set(grayTrackInfo);

spring-cloud-gray-client/src/main/java/cn/springcloud/gray/client/config/GrayClientAutoConfiguration.java

@@ -114,8 +114,8 @@ public GrayClientEnrollInitializingDestroyBean grayClientEnrollInitializingDestr
 
     @Bean
     @ConditionalOnMissingBean
-    public RequestLocalStorage requestLocalStorage() {
-        return new ThreadLocalRequestStorage();
+    public RequestLocalStorage requestLocalStorage(LocalStorageLifeCycle localStorageLifeCycle) {
+        return new ThreadLocalRequestStorage(localStorageLifeCycle);
     }
 
     @Bean

spring-cloud-gray-client/src/main/java/cn/springcloud/gray/request/BaseRequestLocalStorage.java

@@ -0,0 +1,20 @@
+package cn.springcloud.gray.request;
+
+/**
+ * @author saleson
+ * @date 2020-08-22 10:07
+ */
+public abstract class BaseRequestLocalStorage implements RequestLocalStorage {
+
+    private LocalStorageLifeCycle localStorageLifeCycle;
+
+    public BaseRequestLocalStorage(LocalStorageLifeCycle localStorageLifeCycle) {
+        this.localStorageLifeCycle = localStorageLifeCycle;
+    }
+
+    @Override
+    public LocalStorageLifeCycle getLocalStorageLifeCycle() {
+        return localStorageLifeCycle;
+    }
+
+}

spring-cloud-gray-client/src/main/java/cn/springcloud/gray/request/LocalStorageLifeCycle.java

@@ -4,19 +4,33 @@ public interface LocalStorageLifeCycle {
 
     void initContext();
 
+    void initContext(String mark);
+
     void closeContext();
 
+    void closeContext(String mark);
+
 
     public static class NoOpLocalStorageLifeCycle implements LocalStorageLifeCycle {
 
         @Override
         public void initContext() {
-            
+
+        }
+
+        @Override
+        public void initContext(String mark) {
+
         }
 
         @Override
         public void closeContext() {
 
         }
+
+        @Override
+        public void closeContext(String mark) {
+
+        }
     }
 }

spring-cloud-gray-client/src/main/java/cn/springcloud/gray/request/RequestLocalStorage.java

@@ -15,4 +15,13 @@ public interface RequestLocalStorage {
 
     GrayRequest getGrayRequest();
 
+    LocalStorageLifeCycle getLocalStorageLifeCycle();
+
+
+    default void clearAll() {
+        removeGrayRequest();
+        removeGrayTrackInfo();
+        getLocalStorageLifeCycle().closeContext();
+    }
+
 }

spring-cloud-gray-client/src/main/java/cn/springcloud/gray/request/ThreadLocalRequestStorage.java

@@ -1,10 +1,18 @@
 package cn.springcloud.gray.request;
 
-public class ThreadLocalRequestStorage implements RequestLocalStorage {
+public class ThreadLocalRequestStorage extends BaseRequestLocalStorage {
     private ThreadLocal<GrayRequest> grayRequestThreadLocal = new ThreadLocal<>();
     private ThreadLocal<GrayTrackInfo> grayTrackInfoThreadLocal = new ThreadLocal<>();
 
 
+    public ThreadLocalRequestStorage() {
+        this(new LocalStorageLifeCycle.NoOpLocalStorageLifeCycle());
+    }
+
+    public ThreadLocalRequestStorage(LocalStorageLifeCycle localStorageLifeCycle) {
+        super(localStorageLifeCycle);
+    }
+
     @Override
     public void setGrayTrackInfo(GrayTrackInfo grayTrackInfo) {
         grayTrackInfoThreadLocal.set(grayTrackInfo);
@@ -35,4 +43,5 @@ public GrayRequest getGrayRequest() {
         return grayRequestThreadLocal.get();
     }
 
+
 }

spring-cloud-gray-plugins/spring-cloud-gray-plugin-webmvc/src/main/java/cn/springcloud/gray/client/config/GrayMockWebAutoConfiguration.java

@@ -5,14 +5,11 @@
 import cn.springcloud.gray.web.GrayMockFilter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.web.servlet.FilterRegistrationBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import javax.servlet.Filter;
-
 /**
  * @author saleson
  * @date 2020-05-27 11:51
@@ -24,8 +21,7 @@ public class GrayMockWebAutoConfiguration {
 
 
     @Bean
-    @ConditionalOnMissingBean
-    @ConditionalOnClass(Filter.class)
+    @ConditionalOnProperty(value = "gray.mock.enabled", havingValue = "true")
     public GrayMockFilter grayMockFilter(
             MockManager mockManager, RequestLocalStorage requestLocalStorage) {
         return new GrayMockFilter(mockManager, requestLocalStorage);

spring-cloud-gray-plugins/spring-cloud-gray-plugin-webmvc/src/main/java/cn/springcloud/gray/client/config/GrayTrackWebMvcConfiguration.java

@@ -15,7 +15,6 @@
 import org.springframework.boot.web.servlet.FilterRegistrationBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.Import;
 
 import javax.servlet.Filter;
 
@@ -60,7 +59,7 @@ public FilterRegistrationBean grayTraceFilter(GrayTrackFilter filter) {
 
     //    @Configuration
     @ConditionalOnProperty(value = {"gray.hystrix.threadTransmitStrategy"}, havingValue = "HYSTRIX_REQUEST_LOCAL_STORAGE")
-    @Import(HystrixGrayTrackWebConfiguration.class)
+//    @Import(HystrixGrayTrackWebConfiguration.class)
     public static class HystrixRequestLocalStorageConfiguration {
         @Bean
         public RequestLocalStorage requestLocalStorage() {

spring-cloud-gray-plugins/spring-cloud-gray-plugin-webmvc/src/main/java/cn/springcloud/gray/web/GrayTrackFilter.java

@@ -28,13 +28,14 @@ public void init(FilterConfig filterConfig) throws ServletException {
     @Override
     public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
         GrayHttpTrackInfo webTrack = new GrayHttpTrackInfo();
-//        trackors.forEach(trackor -> trackor.call(webTrack, (HttpServletRequest) request));
-        grayTrackHolder.recordGrayTrack(webTrack, new ServletHttpRequestWrapper((HttpServletRequest) request));
-        requestLocalStorage.setGrayTrackInfo(webTrack);
+        requestLocalStorage.getLocalStorageLifeCycle().initContext(GrayTrackFilter.class.getName());
         try {
+            grayTrackHolder.recordGrayTrack(webTrack, new ServletHttpRequestWrapper((HttpServletRequest) request));
+            requestLocalStorage.setGrayTrackInfo(webTrack);
             chain.doFilter(request, response);
         } finally {
             requestLocalStorage.removeGrayTrackInfo();
+            requestLocalStorage.getLocalStorageLifeCycle().closeContext(GrayTrackFilter.class.getName());
         }
     }