Fix re-transform bug when enhanced class proxy parent method (#659)

kylixs · web-flow · commit 4f2f81e18197 · 2023-12-23T00:01:30.000+08:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -8,7 +8,7 @@ Release Notes.
 * Fix NoSuchMethodError in mvc-annotation-commons and change deprecated method.
 * fix forkjoinpool plugin in JDK11。
 * Support for tracing spring-cloud-gateway 4.x in gateway-4.x-plugin.
-
+* Fix re-transform bug when plugin enhanced class proxy parent method.
 
 #### Documentation
 
diff --git a/apm-sniffer/apm-agent/src/main/java/org/apache/skywalking/apm/agent/SkyWalkingAgent.java b/apm-sniffer/apm-agent/src/main/java/org/apache/skywalking/apm/agent/SkyWalkingAgent.java
@@ -32,12 +32,14 @@
 import net.bytebuddy.description.NamedElement;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.scaffold.MethodGraph;
 import net.bytebuddy.dynamic.scaffold.TypeValidation;
 import org.apache.skywalking.apm.agent.bytebuddy.SWAuxiliaryTypeNamingStrategy;
 import net.bytebuddy.implementation.SWImplementationContextFactory;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 import net.bytebuddy.utility.JavaModule;
+import org.apache.skywalking.apm.agent.bytebuddy.SWMethodGraphCompilerDelegate;
 import org.apache.skywalking.apm.agent.bytebuddy.SWMethodNameTransformer;
 import org.apache.skywalking.apm.agent.core.boot.AgentPackageNotFoundException;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
@@ -99,6 +101,23 @@ public static void premain(String agentArgs, Instrumentation instrumentation) th
             return;
         }
 
+        try {
+            installClassTransformer(instrumentation, pluginFinder);
+        } catch (Exception e) {
+            LOGGER.error(e, "Skywalking agent installed class transformer failure.");
+        }
+
+        try {
+            ServiceManager.INSTANCE.boot();
+        } catch (Exception e) {
+            LOGGER.error(e, "Skywalking agent boot failure.");
+        }
+
+        Runtime.getRuntime()
+               .addShutdownHook(new Thread(ServiceManager.INSTANCE::shutdown, "skywalking service shutdown thread"));
+    }
+
+    static void installClassTransformer(Instrumentation instrumentation, PluginFinder pluginFinder) throws Exception {
         LOGGER.info("Skywalking agent begin to install transformer ...");
 
         AgentBuilder agentBuilder = newAgentBuilder().ignore(
@@ -116,15 +135,13 @@ public static void premain(String agentArgs, Instrumentation instrumentation) th
         try {
             agentBuilder = BootstrapInstrumentBoost.inject(pluginFinder, instrumentation, agentBuilder, edgeClasses);
         } catch (Exception e) {
-            LOGGER.error(e, "SkyWalking agent inject bootstrap instrumentation failure. Shutting down.");
-            return;
+            throw new Exception("SkyWalking agent inject bootstrap instrumentation failure. Shutting down.", e);
         }
 
         try {
             agentBuilder = JDK9ModuleExporter.openReadEdge(instrumentation, agentBuilder, edgeClasses);
         } catch (Exception e) {
-            LOGGER.error(e, "SkyWalking agent open read edge in JDK 9+ failure. Shutting down.");
-            return;
+            throw new Exception("SkyWalking agent open read edge in JDK 9+ failure. Shutting down.", e);
         }
 
         agentBuilder.type(pluginFinder.buildMatch())
@@ -137,15 +154,6 @@ public static void premain(String agentArgs, Instrumentation instrumentation) th
         PluginFinder.pluginInitCompleted();
 
         LOGGER.info("Skywalking agent transformer has installed.");
-
-        try {
-            ServiceManager.INSTANCE.boot();
-        } catch (Exception e) {
-            LOGGER.error(e, "Skywalking agent boot failure.");
-        }
-
-        Runtime.getRuntime()
-               .addShutdownHook(new Thread(ServiceManager.INSTANCE::shutdown, "skywalking service shutdown thread"));
     }
 
     /**
@@ -156,7 +164,8 @@ private static AgentBuilder newAgentBuilder() {
         final ByteBuddy byteBuddy = new ByteBuddy()
                 .with(TypeValidation.of(Config.Agent.IS_OPEN_DEBUGGING_CLASS))
                 .with(new SWAuxiliaryTypeNamingStrategy(NAME_TRAIT))
-                .with(new SWImplementationContextFactory(NAME_TRAIT));
+                .with(new SWImplementationContextFactory(NAME_TRAIT))
+                .with(new SWMethodGraphCompilerDelegate(MethodGraph.Compiler.DEFAULT));
 
         return new SWAgentBuilderDefault(byteBuddy, new SWNativeMethodStrategy(NAME_TRAIT))
                 .with(new SWDescriptionStrategy(NAME_TRAIT));
diff --git a/apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-2.x-3.x-plugin/pom.xml b/apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-2.x-3.x-plugin/pom.xml
@@ -42,5 +42,11 @@
             <version>${jedis.version}</version>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.skywalking</groupId>
+            <artifactId>apm-agent</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>
diff --git a/apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-2.x-3.x-plugin/src/test/java/org/apache/skywalking/apm/agent/JedisInstrumentationTest.java b/apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-2.x-3.x-plugin/src/test/java/org/apache/skywalking/apm/agent/JedisInstrumentationTest.java
@@ -0,0 +1,83 @@
+/*
+ *   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.agent;
+
+import com.google.common.base.Stopwatch;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import org.apache.skywalking.apm.agent.core.logging.core.SystemOutWriter;
+import org.apache.skywalking.apm.agent.core.plugin.AbstractClassEnhancePluginDefine;
+import org.apache.skywalking.apm.agent.core.plugin.ByteBuddyCoreClasses;
+import org.apache.skywalking.apm.agent.core.plugin.PluginFinder;
+import org.apache.skywalking.apm.plugin.jedis.v3.define.JedisInstrumentation;
+import org.junit.Assert;
+import org.junit.Test;
+import redis.clients.jedis.Jedis;
+
+import java.lang.instrument.Instrumentation;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+public class JedisInstrumentationTest {
+
+    @Test
+    public void test() throws Exception {
+        // tested plugins
+        List<AbstractClassEnhancePluginDefine> plugins = Arrays.asList(new JedisInstrumentation());
+
+        // remove shade prefix
+        String[] classes = ByteBuddyCoreClasses.CLASSES;
+        for (int i = 0; i < classes.length; i++) {
+            classes[i] = classes[i].replaceFirst("org.apache.skywalking.apm.dependencies.", "");
+        }
+
+        Instrumentation instrumentation = ByteBuddyAgent.install();
+        SkyWalkingAgent.installClassTransformer(instrumentation, new PluginFinder(plugins));
+
+        // first load
+        Jedis jedis = new Jedis();
+        try {
+            jedis.get("mykey");
+        } catch (Exception e) {
+            Assert.assertTrue(e.toString(), e.toString().contains("JedisConnectionException"));
+        }
+
+        log("Do re-transform class : redis.clients.jedis.Jedis ..");
+        Stopwatch stopwatch = Stopwatch.createStarted();
+
+        // re-transform class
+        for (int i = 0; i < 4; i++) {
+            stopwatch.reset();
+            stopwatch.start();
+            instrumentation.retransformClasses(Jedis.class);
+            long elapsed = stopwatch.elapsed(TimeUnit.MILLISECONDS);
+            log("Re-transform class cost: " + elapsed);
+        }
+
+        // test after re-transform class
+        try {
+            jedis.get("mykey");
+        } catch (Exception e) {
+            Assert.assertTrue(e.toString(), e.toString().contains("JedisConnectionException"));
+        }
+    }
+
+    private void log(String message) {
+        SystemOutWriter.INSTANCE.write(message);
+    }
+}
diff --git a/apm-sniffer/bytebuddy-patch/src/main/java/net/bytebuddy/agent/builder/SWDescriptionStrategy.java b/apm-sniffer/bytebuddy-patch/src/main/java/net/bytebuddy/agent/builder/SWDescriptionStrategy.java
@@ -29,7 +29,6 @@
 import net.bytebuddy.pool.TypePool;
 import net.bytebuddy.utility.JavaModule;
 import net.bytebuddy.utility.nullability.MaybeNull;
-import org.apache.skywalking.apm.agent.core.plugin.AbstractClassEnhancePluginDefine;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
 
 import java.io.Serializable;
@@ -38,6 +37,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
 /**
@@ -109,7 +110,7 @@ public TypeDescription apply(String name,
             }
         }
         // wrap result
-        return new SWTypeDescriptionWrapper(delegate.apply(name, type, typePool, circularityLock, classLoader, module), nameTrait);
+        return new SWTypeDescriptionWrapper(delegate.apply(name, type, typePool, circularityLock, classLoader, module), nameTrait, classLoader, name);
     }
 
     /**
@@ -122,25 +123,40 @@ static class SWTypeDescriptionWrapper extends TypeDescription.AbstractBase imple
          */
         private static final long serialVersionUID = 1L;
 
-        private static final List<String> IGNORED_INTERFACES = Arrays.asList(EnhancedInstance.class.getName());
-
-        private static final List<String> IGNORED_FIELDS = Arrays.asList(AbstractClassEnhancePluginDefine.CONTEXT_ATTR_NAME);
+        /**
+         * Original type cache.
+         * classloader hashcode -> ( typeName -> type cache )
+         */
+        private static final Map<Integer, Map<String, TypeCache>> CLASS_LOADER_TYPE_CACHE = new ConcurrentHashMap<>();
 
-        private static final List<String> IGNORED_METHODS = Arrays.asList("getSkyWalkingDynamicField", "setSkyWalkingDynamicField");
+        private static final List<String> IGNORED_INTERFACES = Arrays.asList(EnhancedInstance.class.getName());
 
         private MethodList<MethodDescription.InDefinedShape> methods;
 
         private FieldList<FieldDescription.InDefinedShape> fields;
 
         private final String nameTrait;
 
+        private ClassLoader classLoader;
+
+        private String typeName;
+
         private TypeList.Generic interfaces;
 
         private TypeDescription delegate;
 
-        public SWTypeDescriptionWrapper(TypeDescription delegate, String nameTrait) {
+        public SWTypeDescriptionWrapper(TypeDescription delegate, String nameTrait, ClassLoader classLoader, String typeName) {
             this.delegate = delegate;
             this.nameTrait = nameTrait;
+            this.classLoader = classLoader;
+            this.typeName = typeName;
+        }
+
+        private TypeCache getTypeCache() {
+            int classLoaderHashCode = classLoader != null ? classLoader.hashCode() : 0;
+            Map<String, TypeCache> typeCacheMap = CLASS_LOADER_TYPE_CACHE.computeIfAbsent(classLoaderHashCode, k -> new ConcurrentHashMap<>());
+            TypeCache typeCache = typeCacheMap.computeIfAbsent(typeName, k -> new TypeCache(typeName));
+            return typeCache;
         }
 
         @Override
@@ -164,14 +180,16 @@ public TypeList.Generic getInterfaces() {
         public FieldList<FieldDescription.InDefinedShape> getDeclaredFields() {
             if (this.fields == null) {
                 FieldList<FieldDescription.InDefinedShape> declaredFields = delegate.getDeclaredFields();
-                if (declaredFields.stream()
-                        .anyMatch(f -> f.getName().contains(nameTrait) || IGNORED_FIELDS.contains(f.getName()))) {
-                    // Remove dynamic field tokens generated by SkyWalking
+                TypeCache typeCache = getTypeCache();
+                if (typeCache.fieldNames == null) {
+                    // save origin fields
+                    typeCache.fieldNames = declaredFields.stream().map(WithRuntimeName::getName).collect(Collectors.toSet());
+                    fields = declaredFields;
+                } else {
+                    // return origin fields
                     fields = new FieldList.Explicit<>(declaredFields.stream()
-                            .filter(f -> !f.getName().contains(nameTrait) && !IGNORED_FIELDS.contains(f.getName()))
+                            .filter(f -> typeCache.fieldNames.contains(f.getName()))
                             .collect(Collectors.toList()));
-                } else {
-                    fields = declaredFields;
                 }
             }
             return fields;
@@ -181,14 +199,17 @@ public FieldList<FieldDescription.InDefinedShape> getDeclaredFields() {
         public MethodList<MethodDescription.InDefinedShape> getDeclaredMethods() {
             if (this.methods == null) {
                 MethodList<MethodDescription.InDefinedShape> declaredMethods = delegate.getDeclaredMethods();
-                if (declaredMethods.stream()
-                        .anyMatch(m -> m.getName().contains(nameTrait) || IGNORED_METHODS.contains(m.getName()))) {
-                    // Remove dynamic method tokens generated by SkyWalking
+                TypeCache typeCache = getTypeCache();
+                if (typeCache.methodCodes == null) {
+                    // save original methods
+                    typeCache.methodCodes = declaredMethods.stream().map(m -> m.toString().hashCode()).collect(Collectors.toSet());
+                    methods = declaredMethods;
+                } else {
+                    // return original methods in the same order, remove dynamic method tokens generated by SkyWalking and ByteBuddy
+                    // remove generated methods for delegating superclass methods, such as Jedis.
                     methods = new MethodList.Explicit<>(declaredMethods.stream()
-                            .filter(m -> !m.getName().contains(nameTrait) && !IGNORED_METHODS.contains(m.getName()))
+                            .filter(m -> typeCache.methodCodes.contains(m.toString().hashCode()))
                             .collect(Collectors.toList()));
-                } else {
-                    methods = declaredMethods;
                 }
             }
             return methods;
@@ -332,4 +353,14 @@ public int getModifiers() {
             return delegate.getModifiers();
         }
     }
+
+    static class TypeCache {
+        private String typeName;
+        private Set<Integer> methodCodes;
+        private Set<String> fieldNames;
+
+        public TypeCache(String typeName) {
+            this.typeName = typeName;
+        }
+    }
 }
diff --git a/apm-sniffer/bytebuddy-patch/src/main/java/org/apache/skywalking/apm/agent/bytebuddy/SWMethodGraphCompilerDelegate.java b/apm-sniffer/bytebuddy-patch/src/main/java/org/apache/skywalking/apm/agent/bytebuddy/SWMethodGraphCompilerDelegate.java
