gateway & webflux compatible with these scene (#3419)

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/BodyInserterResponseMethodInterceptor.java → apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/AbstractServerResponseMethodInterceptor.java

@@ -33,7 +33,7 @@ import org.springframework.web.server.ServerWebExchange;
 
 import java.lang.reflect.Method;
 
-public class BodyInserterResponseMethodInterceptor implements InstanceMethodsAroundInterceptor {
+public class AbstractServerResponseMethodInterceptor implements InstanceMethodsAroundInterceptor {
 
     /**
      * The error reason
@@ -43,22 +43,24 @@ public class BodyInserterResponseMethodInterceptor implements InstanceMethodsAro
 
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
-        EnhancedInstance instance = (EnhancedInstance) allArguments[0];
-        AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
-        if (span == null) {
-            return;
-        }
-        ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
-        HttpStatus status = exchange.getResponse().getStatusCode();
-        if (status != null && status.value() >= 400) {
-            span.errorOccurred();
-            if (exchange.getAttribute(ERROR_ATTRIBUTE) != null) {
-                span.log((Throwable) exchange.getAttribute(ERROR_ATTRIBUTE));
+        EnhancedInstance instance = DispatcherHandlerHandleMethodInterceptor.getInstance(allArguments[0]);
+        if (instance != null) {
+            AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
+            if (span == null) {
+                return;
+            }
+            ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
+            HttpStatus status = exchange.getResponse().getStatusCode();
+            if (status != null && status.value() >= 400) {
+                span.errorOccurred();
+                if (exchange.getAttribute(ERROR_ATTRIBUTE) != null) {
+                    span.log((Throwable) exchange.getAttribute(ERROR_ATTRIBUTE));
+                }
+                Tags.STATUS_CODE.set(span, Integer.toString(status.value()));
             }
-            Tags.STATUS_CODE.set(span, Integer.toString(status.value()));
+            ContextManager.stopSpan(span);
+            ((EnhancedInstance) allArguments[0]).setSkyWalkingDynamicField(null);
         }
-        ContextManager.stopSpan(span);
-        ((EnhancedInstance) allArguments[0]).setSkyWalkingDynamicField(null);
     }
 
     @Override

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/DispatcherHandlerHandleMethodInterceptor.java

@@ -34,6 +34,8 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInt
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 import org.springframework.http.HttpHeaders;
 import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.ServerWebExchangeDecorator;
+import org.springframework.web.server.adapter.DefaultServerWebExchange;
 
 import java.lang.reflect.Method;
 import java.util.List;
@@ -45,23 +47,25 @@ public class DispatcherHandlerHandleMethodInterceptor implements InstanceMethods
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        ContextCarrier contextCarrier = new ContextCarrier();
-        CarrierItem next = contextCarrier.items();
-        ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
-        EnhancedInstance instance = (EnhancedInstance) allArguments[0];
-        HttpHeaders headers = exchange.getRequest().getHeaders();
-        while (next.hasNext()) {
-            next = next.next();
-            List<String> header = headers.get(next.getHeadKey());
-            if (header != null && header.size() > 0) {
-                next.setHeadValue(header.get(0));
+        EnhancedInstance instance = DispatcherHandlerHandleMethodInterceptor.getInstance(allArguments[0]);
+        if (instance != null) {
+            ContextCarrier contextCarrier = new ContextCarrier();
+            CarrierItem next = contextCarrier.items();
+            ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
+            HttpHeaders headers = exchange.getRequest().getHeaders();
+            while (next.hasNext()) {
+                next = next.next();
+                List<String> header = headers.get(next.getHeadKey());
+                if (header != null && header.size() > 0) {
+                    next.setHeadValue(header.get(0));
+                }
             }
+            AbstractSpan span = ContextManager.createEntrySpan(WIP_OPERATION_NAME, contextCarrier);
+            span.setComponent(ComponentsDefine.SPRING_WEBFLUX);
+            SpanLayer.asHttp(span);
+            Tags.URL.set(span, exchange.getRequest().getURI().toString());
+            instance.setSkyWalkingDynamicField(span);
         }
-        AbstractSpan span = ContextManager.createEntrySpan(WIP_OPERATION_NAME, contextCarrier);
-        span.setComponent(ComponentsDefine.SPRING_WEBFLUX);
-        SpanLayer.asHttp(span);
-        Tags.URL.set(span, exchange.getRequest().getURI().toString());
-        instance.setSkyWalkingDynamicField(span);
     }
 
     @Override
@@ -74,4 +78,17 @@ public class DispatcherHandlerHandleMethodInterceptor implements InstanceMethods
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                       Class<?>[] argumentsTypes, Throwable t) {
     }
+
+    public static EnhancedInstance getInstance(Object o) {
+        EnhancedInstance instance = null;
+        if (o instanceof ServerWebExchangeDecorator) {
+            ServerWebExchange delegate = ((ServerWebExchangeDecorator) o).getDelegate();
+            if (delegate instanceof DefaultServerWebExchange) {
+                instance = (EnhancedInstance) delegate;
+            }
+        } else if (o instanceof DefaultServerWebExchange) {
+            instance = (EnhancedInstance) o;
+        }
+        return instance;
+    }
 }

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/DispatcherHandlerHandleResultMethodInterceptor.java

@@ -40,11 +40,13 @@ public class DispatcherHandlerHandleResultMethodInterceptor implements InstanceM
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                               Object ret) throws Throwable {
-        EnhancedInstance instance = (EnhancedInstance) allArguments[0];
-        AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
-        if (span != null) {
-            ContextManager.stopSpan(span);
-            instance.setSkyWalkingDynamicField(null);
+        EnhancedInstance instance = DispatcherHandlerHandleMethodInterceptor.getInstance(allArguments[0]);
+        if (instance != null) {
+            AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
+            if (span != null) {
+                ContextManager.stopSpan(span);
+                instance.setSkyWalkingDynamicField(null);
+            }
         }
         return ret;
     }

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/DispatcherHandlerInvokeHandlerMethodInterceptor.java

@@ -41,25 +41,27 @@ public class DispatcherHandlerInvokeHandlerMethodInterceptor implements Instance
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        EnhancedInstance instance = (EnhancedInstance) allArguments[0];
-        AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
-        if (span == null) {
-            return;
-        }
-        String handleClassName = allArguments[1].getClass().getSimpleName();
-        int index = handleClassName.indexOf(ROUTER_SEARCH);
-        if (index != -1) {
-            String operationName = handleClassName.substring(0, index);
-            try {
-                Field field = allArguments[1].getClass().getDeclaredField(ROUTER_FIELD);
-                field.setAccessible(true);
-                operationName = operationName + DOT + field.get(allArguments[1]).getClass().getName();
-            } catch (NoSuchFieldException ignore) {
+        EnhancedInstance instance = DispatcherHandlerHandleMethodInterceptor.getInstance(allArguments[0]);
+        if (instance != null) {
+            AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
+            if (span == null) {
+                return;
+            }
+            String handleClassName = allArguments[1].getClass().getSimpleName();
+            int index = handleClassName.indexOf(ROUTER_SEARCH);
+            if (index != -1) {
+                String operationName = handleClassName.substring(0, index);
+                try {
+                    Field field = allArguments[1].getClass().getDeclaredField(ROUTER_FIELD);
+                    field.setAccessible(true);
+                    operationName = operationName + DOT + field.get(allArguments[1]).getClass().getName();
+                } catch (NoSuchFieldException ignore) {
+                }
+                span.setOperationName(operationName);
+            } else if (allArguments[1] instanceof HandlerMethod) {
+                HandlerMethod handler = (HandlerMethod) allArguments[1];
+                span.setOperationName(getHandlerMethodOperationName(handler));
             }
-            span.setOperationName(operationName);
-        } else if (allArguments[1] instanceof HandlerMethod) {
-            HandlerMethod handler = (HandlerMethod) allArguments[1];
-            span.setOperationName(getHandlerMethodOperationName(handler));
         }
     }
 

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/DefaultServerWebExchangeConstructorInterceptor.java → apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/ServerWebExchangeConstructorInterceptor.java

@@ -25,7 +25,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceC
  *
  * @author zhaoyuguang
  */
-public class DefaultServerWebExchangeConstructorInterceptor implements InstanceConstructorInterceptor {
+public class ServerWebExchangeConstructorInterceptor implements InstanceConstructorInterceptor {
     @Override
     public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
     }

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/define/BodyInserterResponseInstrumentation.java → apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/define/AbstractServerResponseInstrumentation.java

@@ -32,12 +32,13 @@ import static org.apache.skywalking.apm.agent.core.plugin.match.MultiClassNameMa
  * @author zhaoyuguang
  */
 
-public class BodyInserterResponseInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
+public class AbstractServerResponseInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
 
     @Override
     protected ClassMatch enhanceClass() {
         return byMultiClassMatch("org.springframework.web.reactive.function.server.DefaultServerResponseBuilder$BodyInserterResponse",
-                "org.springframework.web.reactive.function.server.DefaultServerResponseBuilder$BodyInserterServerResponse");
+                "org.springframework.web.reactive.function.server.DefaultServerResponseBuilder$BodyInserterServerResponse",
+                "org.springframework.web.reactive.function.server.DefaultEntityResponseBuilder$DefaultEntityResponse");
     }
 
     @Override
@@ -54,7 +55,7 @@ public class BodyInserterResponseInstrumentation extends ClassInstanceMethodsEnh
                 }
 
                 @Override public String getMethodsInterceptor() {
-                    return "org.apache.skywalking.apm.plugin.spring.webflux.v5.BodyInserterResponseMethodInterceptor";
+                    return "org.apache.skywalking.apm.plugin.spring.webflux.v5.AbstractServerResponseMethodInterceptor";
                 }
 
                 @Override public boolean isOverrideArgs() {

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/define/DefaultServerWebExchangeInstrumentation.java → apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/webflux/v5/define/ServerWebExchangeInstrumentation.java

@@ -26,13 +26,13 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInst
 import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
 
 import static net.bytebuddy.matcher.ElementMatchers.any;
-import static org.apache.skywalking.apm.agent.core.plugin.match.MultiClassNameMatch.byMultiClassMatch;
+import static org.apache.skywalking.apm.agent.core.plugin.match.NameMatch.byName;
 
 /**
  * @author zhaoyuguang
  */
 
-public class DefaultServerWebExchangeInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
+public class ServerWebExchangeInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
     @Override
     public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
         return new ConstructorInterceptPoint[]{
@@ -44,7 +44,7 @@ public class DefaultServerWebExchangeInstrumentation extends ClassInstanceMethod
 
                 @Override
                 public String getConstructorInterceptor() {
-                    return "org.apache.skywalking.apm.plugin.spring.webflux.v5.DefaultServerWebExchangeConstructorInterceptor";
+                    return "org.apache.skywalking.apm.plugin.spring.webflux.v5.ServerWebExchangeConstructorInterceptor";
                 }
             }
         };
@@ -57,6 +57,6 @@ public class DefaultServerWebExchangeInstrumentation extends ClassInstanceMethod
 
     @Override
     protected ClassMatch enhanceClass() {
-        return byMultiClassMatch("org.springframework.web.server.adapter.DefaultServerWebExchange", "org.springframework.web.server.ServerWebExchangeDecorator");
+        return byName("org.springframework.web.server.adapter.DefaultServerWebExchange");
     }
 }

apm-sniffer/apm-sdk-plugin/spring-plugins/webflux-5.x-plugin/src/main/resources/skywalking-plugin.def

@@ -15,5 +15,5 @@
 # limitations under the License.
 
 spring-webflux-5.x=org.apache.skywalking.apm.plugin.spring.webflux.v5.define.DispatcherHandlerInstrumentation
-spring-webflux-5.x=org.apache.skywalking.apm.plugin.spring.webflux.v5.define.DefaultServerWebExchangeInstrumentation
-spring-webflux-5.x=org.apache.skywalking.apm.plugin.spring.webflux.v5.define.BodyInserterResponseInstrumentation
\ No newline at end of file
+spring-webflux-5.x=org.apache.skywalking.apm.plugin.spring.webflux.v5.define.ServerWebExchangeInstrumentation
+spring-webflux-5.x=org.apache.skywalking.apm.plugin.spring.webflux.v5.define.AbstractServerResponseInstrumentation
\ No newline at end of file

apm-sniffer/optional-plugins/optional-spring-plugins/optional-spring-cloud/gateway-2.1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/cloud/gateway/v21x/NettyRoutingFilterInterceptor.java

@@ -27,6 +27,8 @@ import org.apache.skywalking.apm.plugin.spring.cloud.gateway.v21x.context.Consta
 import org.apache.skywalking.apm.plugin.spring.cloud.gateway.v21x.context.SWTransmitter;
 import org.springframework.cloud.gateway.route.Route;
 import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.ServerWebExchangeDecorator;
+import org.springframework.web.server.adapter.DefaultServerWebExchange;
 
 import java.lang.reflect.Method;
 
@@ -43,16 +45,19 @@ public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInter
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
-        AbstractSpan span = (AbstractSpan) ((EnhancedInstance) allArguments[0]).getSkyWalkingDynamicField();
-        String operationName = SPRING_CLOUD_GATEWAY_ROUTE_PREFIX;
-        if (span != null) {
-            Route route = exchange.getRequiredAttribute(GATEWAY_ROUTE_ATTR);
-            operationName = operationName + route.getId();
-            span.setOperationName(operationName);
-            SWTransmitter transmitter = new SWTransmitter(span.prepareForAsync(), ContextManager.capture(), operationName);
-            ContextManager.stopSpan(span);
-            ContextManager.getRuntimeContext().put(Constants.SPRING_CLOUD_GATEWAY_TRANSMITTER, transmitter);
+        EnhancedInstance instance = NettyRoutingFilterInterceptor.getInstance(allArguments[0]);
+        if (instance != null) {
+            ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
+            AbstractSpan span = (AbstractSpan) instance.getSkyWalkingDynamicField();
+            String operationName = SPRING_CLOUD_GATEWAY_ROUTE_PREFIX;
+            if (span != null) {
+                Route route = exchange.getRequiredAttribute(GATEWAY_ROUTE_ATTR);
+                operationName = operationName + route.getId();
+                span.setOperationName(operationName);
+                SWTransmitter transmitter = new SWTransmitter(span.prepareForAsync(), ContextManager.capture(), operationName);
+                ContextManager.stopSpan(span);
+                ContextManager.getRuntimeContext().put(Constants.SPRING_CLOUD_GATEWAY_TRANSMITTER, transmitter);
+            }
         }
     }
 
@@ -70,4 +75,17 @@ public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInter
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                       Class<?>[] argumentsTypes, Throwable t) {
     }
+
+    public static EnhancedInstance getInstance(Object o) {
+        EnhancedInstance instance = null;
+        if (o instanceof ServerWebExchangeDecorator) {
+            ServerWebExchange delegate = ((ServerWebExchangeDecorator) o).getDelegate();
+            if (delegate instanceof DefaultServerWebExchange) {
+                instance = (EnhancedInstance) delegate;
+            }
+        } else if (o instanceof DefaultServerWebExchange) {
+            instance = (EnhancedInstance) o;
+        }
+        return instance;
+    }
 }