增加TraceId传递功能

discovery-common/src/main/java/com/nepxion/discovery/common/constant/DiscoveryConstant.java

@@ -51,4 +51,5 @@ public class DiscoveryConstant {
     public static final String NO = "NO";
     public static final String UNKNOWN = "UNKNOWN";
     public static final String EXT = "ext";
+    public static final String TRACE_ID = "traceId";
 }
\ No newline at end of file

discovery-plugin-strategy-starter-service/src/main/java/com/nepxion/discovery/plugin/strategy/service/aop/FeignStrategyInterceptor.java

@@ -16,18 +16,24 @@ import java.util.Enumeration;
 
 import javax.servlet.http.HttpServletRequest;
 
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.web.context.request.ServletRequestAttributes;
 
+import com.nepxion.discovery.common.constant.DiscoveryConstant;
 import com.nepxion.discovery.plugin.strategy.service.context.ServiceStrategyContextHolder;
+import com.nepxion.discovery.plugin.strategy.service.trace.TraceIdGenerator;
 
 public class FeignStrategyInterceptor implements RequestInterceptor {
     private static final Logger LOG = LoggerFactory.getLogger(FeignStrategyInterceptor.class);
 
     private String requestHeaders;
 
+    @Autowired(required = false)
+    private TraceIdGenerator traceIdGenerator;
+
     @Autowired
     private ServiceStrategyContextHolder serviceStrategyContextHolder;
 
@@ -52,6 +58,18 @@ public class FeignStrategyInterceptor implements RequestInterceptor {
             return;
         }
 
+        String traceId = previousRequest.getHeader(DiscoveryConstant.TRACE_ID);
+        if (StringUtils.isEmpty(traceId) && traceIdGenerator != null) {
+            try {
+                traceId = traceIdGenerator.generate();
+            } catch (Exception e) {
+                LOG.error("Generate trace id failed, ignore to set trace id", e);
+            }
+        }
+        if (StringUtils.isNotEmpty(traceId)) {
+            requestTemplate.header(DiscoveryConstant.TRACE_ID, traceId);
+        }
+
         while (headerNames.hasMoreElements()) {
             String headerName = headerNames.nextElement();
             String header = previousRequest.getHeader(headerName);

discovery-plugin-strategy-starter-service/src/main/java/com/nepxion/discovery/plugin/strategy/service/aop/RestTemplateStrategyInterceptor.java

@@ -14,6 +14,7 @@ import java.util.Enumeration;
 
 import javax.servlet.http.HttpServletRequest;
 
+import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -24,13 +25,18 @@ import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.web.context.request.ServletRequestAttributes;
 
+import com.nepxion.discovery.common.constant.DiscoveryConstant;
 import com.nepxion.discovery.plugin.strategy.service.context.ServiceStrategyContextHolder;
+import com.nepxion.discovery.plugin.strategy.service.trace.TraceIdGenerator;
 
 public class RestTemplateStrategyInterceptor implements ClientHttpRequestInterceptor {
     private static final Logger LOG = LoggerFactory.getLogger(RestTemplateStrategyInterceptor.class);
 
     private String requestHeaders;
 
+    @Autowired(required = false)
+    private TraceIdGenerator traceIdGenerator;
+
     @Autowired
     private ServiceStrategyContextHolder serviceStrategyContextHolder;
 
@@ -56,6 +62,19 @@ public class RestTemplateStrategyInterceptor implements ClientHttpRequestInterce
         }
 
         HttpHeaders headers = request.getHeaders();
+
+        String traceId = previousRequest.getHeader(DiscoveryConstant.TRACE_ID);
+        if (StringUtils.isEmpty(traceId) && traceIdGenerator != null) {
+            try {
+                traceId = traceIdGenerator.generate();
+            } catch (Exception e) {
+                LOG.error("Generate trace id failed, ignore to set trace id", e);
+            }
+        }
+        if (StringUtils.isNotEmpty(traceId)) {
+            headers.add(DiscoveryConstant.TRACE_ID, traceId);
+        }
+
         while (headerNames.hasMoreElements()) {
             String headerName = headerNames.nextElement();
             String header = previousRequest.getHeader(headerName);

discovery-plugin-strategy-starter-service/src/main/java/com/nepxion/discovery/plugin/strategy/service/trace/TraceIdGenerator.java

+package com.nepxion.discovery.plugin.strategy.service.trace;
+
+/**
+ * <p>Title: Nepxion Discovery</p>
+ * <p>Description: Nepxion Discovery</p>
+ * <p>Copyright: Copyright (c) 2017-2050</p>
+ * <p>Company: Nepxion</p>
+ * @author Haojun Ren
+ * @version 1.0
+ */
+
+public interface TraceIdGenerator {
+    String generate();
+}
\ No newline at end of file

discovery-springcloud-example-service/src/main/java/com/nepxion/discovery/plugin/example/service/DiscoveryApplicationA1.java

@@ -28,6 +28,7 @@ import com.nepxion.discovery.plugin.example.service.impl.MyRegisterListener;
 import com.nepxion.discovery.plugin.example.service.impl.MySentinelExceptionHandler;
 import com.nepxion.discovery.plugin.example.service.impl.MySentinelFlowRuleParser;
 import com.nepxion.discovery.plugin.example.service.impl.MySubscriber;
+import com.nepxion.discovery.plugin.example.service.impl.MyTraceIdGenerator;
 import com.nepxion.discovery.plugin.strategy.service.aop.RestTemplateStrategyInterceptor;
 
 @SpringBootApplication
@@ -85,4 +86,9 @@ public class DiscoveryApplicationA1 {
     public MyDiscoveryEnabledStrategy myDiscoveryEnabledStrategy() {
         return new MyDiscoveryEnabledStrategy();
     }
+
+    @Bean
+    public MyTraceIdGenerator myTraceIdGenerator() {
+        return new MyTraceIdGenerator();
+    }
 }
\ No newline at end of file

discovery-springcloud-example-service/src/main/java/com/nepxion/discovery/plugin/example/service/feign/AbstractFeignImpl.java

@@ -11,13 +11,19 @@ package com.nepxion.discovery.plugin.example.service.feign;
 
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.web.context.request.ServletRequestAttributes;
 
+import com.nepxion.discovery.common.constant.DiscoveryConstant;
 import com.nepxion.discovery.plugin.framework.adapter.PluginAdapter;
+import com.nepxion.discovery.plugin.strategy.service.context.ServiceStrategyContextHolder;
 
 public class AbstractFeignImpl {
     @Autowired
     private PluginAdapter pluginAdapter;
 
+    @Autowired
+    private ServiceStrategyContextHolder serviceStrategyContextHolder;
+
     public String doInvoke(String value) {
         String serviceId = pluginAdapter.getServiceId();
         String host = pluginAdapter.getHost();
@@ -25,6 +31,12 @@ public class AbstractFeignImpl {
         String version = pluginAdapter.getVersion();
         String region = pluginAdapter.getRegion();
 
+        String traceId = null;
+        ServletRequestAttributes attributes = serviceStrategyContextHolder.getRestAttributes();
+        if (attributes != null) {
+            traceId = attributes.getRequest().getHeader(DiscoveryConstant.TRACE_ID);
+        }
+
         StringBuilder stringBuilder = new StringBuilder();
         stringBuilder.append(value + " -> " + serviceId);
         stringBuilder.append("[" + host + ":" + port + "]");
@@ -34,6 +46,9 @@ public class AbstractFeignImpl {
         if (StringUtils.isNotEmpty(region)) {
             stringBuilder.append("[Region=" + region + "]");
         }
+        if (StringUtils.isNotEmpty(traceId)) {
+            stringBuilder.append("[TraceId=" + traceId + "]");
+        }
 
         return stringBuilder.toString();
     }

discovery-springcloud-example-service/src/main/java/com/nepxion/discovery/plugin/example/service/impl/MyTraceIdGenerator.java

+package com.nepxion.discovery.plugin.example.service.impl;
+
+/**
+ * <p>Title: Nepxion Discovery</p>
+ * <p>Description: Nepxion Discovery</p>
+ * <p>Copyright: Copyright (c) 2017-2050</p>
+ * <p>Company: Nepxion</p>
+ * @author Haojun Ren
+ * @version 1.0
+ */
+
+import java.util.UUID;
+
+import com.nepxion.discovery.plugin.strategy.service.trace.TraceIdGenerator;
+
+public class MyTraceIdGenerator implements TraceIdGenerator {
+    @Override
+    public String generate() {
+        return UUID.randomUUID().toString();
+    }
+}
\ No newline at end of file

discovery-springcloud-example-service/src/main/java/com/nepxion/discovery/plugin/example/service/rest/AbstractRestImpl.java

@@ -11,13 +11,19 @@ package com.nepxion.discovery.plugin.example.service.rest;
 
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.web.context.request.ServletRequestAttributes;
 
+import com.nepxion.discovery.common.constant.DiscoveryConstant;
 import com.nepxion.discovery.plugin.framework.adapter.PluginAdapter;
+import com.nepxion.discovery.plugin.strategy.service.context.ServiceStrategyContextHolder;
 
 public class AbstractRestImpl {
     @Autowired
     private PluginAdapter pluginAdapter;
 
+    @Autowired
+    private ServiceStrategyContextHolder serviceStrategyContextHolder;
+
     public String doRest(String value) {
         String serviceId = pluginAdapter.getServiceId();
         String host = pluginAdapter.getHost();
@@ -25,6 +31,12 @@ public class AbstractRestImpl {
         String version = pluginAdapter.getVersion();
         String region = pluginAdapter.getRegion();
 
+        String traceId = null;
+        ServletRequestAttributes attributes = serviceStrategyContextHolder.getRestAttributes();
+        if (attributes != null) {
+            traceId = attributes.getRequest().getHeader(DiscoveryConstant.TRACE_ID);
+        }
+
         StringBuilder stringBuilder = new StringBuilder();
         stringBuilder.append(value + " -> " + serviceId);
         stringBuilder.append("[" + host + ":" + port + "]");
@@ -34,6 +46,9 @@ public class AbstractRestImpl {
         if (StringUtils.isNotEmpty(region)) {
             stringBuilder.append("[Region=" + region + "]");
         }
+        if (StringUtils.isNotEmpty(traceId)) {
+            stringBuilder.append("[TraceId=" + traceId + "]");
+        }
 
         return stringBuilder.toString();
     }