Fix the bug that maybe causing memory leak and repeated traceId when use gateway-2.1.x-plugin (#133)

CHANGES.md

@@ -34,6 +34,7 @@ Release Notes.
   instance_properties_json includes these two keys, they would be overrided by the agent core.
 * [Breaking Change] Remove the namespace from `cross process propagation` key.
 * Make sure the parent endpoint in tracing context from existing first ENTRY span, rather than first span only.
+* Fix the bug that maybe causing memory leak and repeated traceId when use gateway-2.1.x-plugin or gateway-3.x-plugin.
 
 #### Documentation
 

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

@@ -31,10 +31,18 @@ import org.springframework.web.server.adapter.DefaultServerWebExchange;
 import static org.apache.skywalking.apm.network.trace.component.ComponentsDefine.SPRING_CLOUD_GATEWAY;
 
 public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInterceptor {
+    private static final String NETTY_ROUTING_FILTER_TRACED_ATTR = NettyRoutingFilterInterceptor.class.getName() + ".isTraced";
+
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
         ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
+        if (isTraced(exchange)) {
+            return;
+        }
+
+        setTracedStatus(exchange);
+
         EnhancedInstance enhancedInstance = getInstance(exchange);
 
         AbstractSpan span = ContextManager.createLocalSpan("SpringCloudGateway/RoutingFilter");
@@ -44,6 +52,14 @@ public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInter
         span.setComponent(SPRING_CLOUD_GATEWAY);
     }
 
+    private static void setTracedStatus(ServerWebExchange exchange) {
+        exchange.getAttributes().put(NETTY_ROUTING_FILTER_TRACED_ATTR, true);
+    }
+
+    private static boolean isTraced(ServerWebExchange exchange) {
+        return exchange.getAttributeOrDefault(NETTY_ROUTING_FILTER_TRACED_ATTR, false);
+    }
+
     public static EnhancedInstance getInstance(Object o) {
         EnhancedInstance instance = null;
         if (o instanceof DefaultServerWebExchange) {

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

@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.apm.plugin.spring.cloud.gateway.v21x;
+
+import java.security.Principal;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.test.tools.AgentServiceRule;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStorage;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStoragePoint;
+import org.apache.skywalking.apm.agent.test.tools.TracingSegmentRunner;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.modules.junit4.PowerMockRunnerDelegate;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.i18n.LocaleContext;
+import org.springframework.http.codec.multipart.Part;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebSession;
+import reactor.core.publisher.Mono;
+
+@RunWith(PowerMockRunner.class)
+@PowerMockRunnerDelegate(TracingSegmentRunner.class)
+public class NettyRoutingFilterInterceptorTest {
+    private final NettyRoutingFilterInterceptor interceptor = new NettyRoutingFilterInterceptor();
+    @Rule
+    public AgentServiceRule serviceRule = new AgentServiceRule();
+    @SegmentStoragePoint
+    private SegmentStorage segmentStorage;
+
+    @Before
+    public void setUp() throws Exception {
+    }
+
+    private static final String NETTY_ROUTING_FILTER_TRACED_ATTR =
+            NettyRoutingFilterInterceptor.class.getName() + ".isTraced";
+
+    private final ServerWebExchange exchange = new ServerWebExchange() {
+        Map<String, Object> attributes = new HashMap<>();
+        @Override
+        public ServerHttpRequest getRequest() {
+            return null;
+        }
+
+        @Override
+        public ServerHttpResponse getResponse() {
+            return null;
+        }
+
+        @Override
+        public Map<String, Object> getAttributes() {
+            return attributes;
+        }
+
+        @Override
+        public Mono<WebSession> getSession() {
+            return null;
+        }
+
+        @Override
+        public <T extends Principal> Mono<T> getPrincipal() {
+            return null;
+        }
+
+        @Override
+        public Mono<MultiValueMap<String, String>> getFormData() {
+            return null;
+        }
+
+        @Override
+        public Mono<MultiValueMap<String, Part>> getMultipartData() {
+            return null;
+        }
+
+        @Override
+        public LocaleContext getLocaleContext() {
+            return null;
+        }
+
+        @Override
+        public ApplicationContext getApplicationContext() {
+            return null;
+        }
+
+        @Override
+        public boolean isNotModified() {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(Instant instant) {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(String s) {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(String s, Instant instant) {
+            return false;
+        }
+
+        @Override
+        public String transformUrl(String s) {
+            return null;
+        }
+
+        @Override
+        public void addUrlTransformer(Function<String, String> function) {
+
+        }
+
+        @Override
+        public String getLogPrefix() {
+            return null;
+        }
+    };
+
+    @Test
+    public void testIsTraced() throws Throwable {
+        interceptor.beforeMethod(null, null, new Object[]{exchange}, null, null);
+        interceptor.afterMethod(null, null, null, null, null);
+        Assert.assertEquals(exchange.getAttributes().get(NETTY_ROUTING_FILTER_TRACED_ATTR), true);
+        Assert.assertNotNull(ContextManager.activeSpan());
+
+        ContextManager.stopSpan();
+
+        interceptor.beforeMethod(null, null, new Object[]{exchange}, null, null);
+        interceptor.afterMethod(null, null, null, null, null);
+        Assert.assertEquals(exchange.getAttributes().get(NETTY_ROUTING_FILTER_TRACED_ATTR), true);
+    }
+}

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

@@ -38,10 +38,18 @@ import static org.apache.skywalking.apm.network.trace.component.ComponentsDefine
  * </p>
  */
 public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInterceptor {
+    private static final String NETTY_ROUTING_FILTER_TRACED_ATTR = NettyRoutingFilterInterceptor.class.getName() + ".isTraced";
 
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
             MethodInterceptResult result) throws Throwable {
+        ServerWebExchange exchange = (ServerWebExchange) allArguments[0];
+        if (isTraced(exchange)) {
+            return;
+        }
+
+        setTracedStatus(exchange);
+
         EnhancedInstance enhancedInstance = getInstance(allArguments[0]);
 
         AbstractSpan span = ContextManager.createLocalSpan("SpringCloudGateway/RoutingFilter");
@@ -51,6 +59,14 @@ public class NettyRoutingFilterInterceptor implements InstanceMethodsAroundInter
         span.setComponent(SPRING_CLOUD_GATEWAY);
     }
 
+    private static void setTracedStatus(ServerWebExchange exchange) {
+        exchange.getAttributes().put(NETTY_ROUTING_FILTER_TRACED_ATTR, true);
+    }
+
+    private static boolean isTraced(ServerWebExchange exchange) {
+        return exchange.getAttributeOrDefault(NETTY_ROUTING_FILTER_TRACED_ATTR, false);
+    }
+
     private EnhancedInstance getInstance(Object o) {
         EnhancedInstance instance = null;
         if (o instanceof EnhancedInstance) {

apm-sniffer/optional-plugins/optional-spring-plugins/optional-spring-cloud/gateway-3.x-plugin/src/test/java/org/apache/skywalking/apm/plugin/spring/cloud/gateway/v3x/NettyRoutingFilterInterceptorTest.java

@@ -18,6 +18,11 @@
 
 package org.apache.skywalking.apm.plugin.spring.cloud.gateway.v3x;
 
+import java.security.Principal;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
 import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
@@ -41,13 +46,22 @@ import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.modules.junit4.PowerMockRunnerDelegate;
 
 import java.util.List;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.i18n.LocaleContext;
+import org.springframework.http.codec.multipart.Part;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebSession;
+import reactor.core.publisher.Mono;
 
 @RunWith(PowerMockRunner.class)
 @PowerMockRunnerDelegate(TracingSegmentRunner.class)
 public class NettyRoutingFilterInterceptorTest {
-
-    private final EnhancedInstance enhancedInstance = new EnhancedInstance() {
+    private static class ServerWebExchangeEnhancedInstance implements ServerWebExchange, EnhancedInstance {
         private ContextSnapshot snapshot;
+        Map<String, Object> attributes = new HashMap<>();
 
         @Override
         public Object getSkyWalkingDynamicField() {
@@ -58,7 +72,89 @@ public class NettyRoutingFilterInterceptorTest {
         public void setSkyWalkingDynamicField(Object value) {
             this.snapshot = (ContextSnapshot) value;
         }
-    };
+
+        @Override
+        public ServerHttpRequest getRequest() {
+            return null;
+        }
+
+        @Override
+        public ServerHttpResponse getResponse() {
+            return null;
+        }
+
+        @Override
+        public Map<String, Object> getAttributes() {
+            return attributes;
+        }
+
+        @Override
+        public Mono<WebSession> getSession() {
+            return null;
+        }
+
+        @Override
+        public <T extends Principal> Mono<T> getPrincipal() {
+            return null;
+        }
+
+        @Override
+        public Mono<MultiValueMap<String, String>> getFormData() {
+            return null;
+        }
+
+        @Override
+        public Mono<MultiValueMap<String, Part>> getMultipartData() {
+            return null;
+        }
+
+        @Override
+        public LocaleContext getLocaleContext() {
+            return null;
+        }
+
+        @Override
+        public ApplicationContext getApplicationContext() {
+            return null;
+        }
+
+        @Override
+        public boolean isNotModified() {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(Instant instant) {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(String s) {
+            return false;
+        }
+
+        @Override
+        public boolean checkNotModified(String s, Instant instant) {
+            return false;
+        }
+
+        @Override
+        public String transformUrl(String s) {
+            return null;
+        }
+
+        @Override
+        public void addUrlTransformer(Function<String, String> function) {
+
+        }
+
+        @Override
+        public String getLogPrefix() {
+            return null;
+        }
+    }
+
+    private final ServerWebExchangeEnhancedInstance enhancedInstance = new ServerWebExchangeEnhancedInstance();
     private final NettyRoutingFilterInterceptor interceptor = new NettyRoutingFilterInterceptor();
     @Rule
     public AgentServiceRule serviceRule = new AgentServiceRule();
@@ -101,4 +197,21 @@ public class NettyRoutingFilterInterceptorTest {
         SpanAssert.assertComponent(spans.get(1), ComponentsDefine.SPRING_WEBFLUX);
         SpanAssert.assertLayer(spans.get(1), SpanLayer.HTTP);
     }
+
+    private static final String NETTY_ROUTING_FILTER_TRACED_ATTR =
+            NettyRoutingFilterInterceptor.class.getName() + ".isTraced";
+
+    @Test
+    public void testIsTraced() throws Throwable {
+        interceptor.beforeMethod(null, null, new Object[]{enhancedInstance}, null, null);
+        interceptor.afterMethod(null, null, null, null, null);
+        Assert.assertEquals(enhancedInstance.getAttributes().get(NETTY_ROUTING_FILTER_TRACED_ATTR), true);
+        Assert.assertNotNull(ContextManager.activeSpan());
+
+        ContextManager.stopSpan();
+
+        interceptor.beforeMethod(null, null, new Object[]{enhancedInstance}, null, null);
+        interceptor.afterMethod(null, null, null, null, null);
+        Assert.assertEquals(enhancedInstance.getAttributes().get(NETTY_ROUTING_FILTER_TRACED_ATTR), true);
+    }
 }