Merge pull request #478 from wu-sheng/feature/multi-enhance

Support multi plugin defines targeting the same class. 

apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/plugin/AbstractClassEnhancePluginDefine.java

@@ -26,7 +26,7 @@ public abstract class AbstractClassEnhancePluginDefine {
      * @throws PluginException, when set builder failure.
      */
     public DynamicType.Builder<?> define(String transformClassName,
-        DynamicType.Builder<?> builder, ClassLoader classLoader) throws PluginException {
+        DynamicType.Builder<?> builder, ClassLoader classLoader, EnhanceContext context) throws PluginException {
         String interceptorDefineClassName = this.getClass().getName();
 
         if (StringUtil.isEmpty(transformClassName)) {
@@ -53,15 +53,16 @@ public abstract class AbstractClassEnhancePluginDefine {
         /**
          * find origin class source code for interceptor
          */
-        DynamicType.Builder<?> newClassBuilder = this.enhance(transformClassName, builder, classLoader);
+        DynamicType.Builder<?> newClassBuilder = this.enhance(transformClassName, builder, classLoader, context);
 
+        context.initializationStageCompleted();
         logger.debug("enhance class {} by {} completely.", transformClassName, interceptorDefineClassName);
 
         return newClassBuilder;
     }
 
     protected abstract DynamicType.Builder<?> enhance(String enhanceOriginClassName,
-        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader) throws PluginException;
+        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader, EnhanceContext context) throws PluginException;
 
     /**
      * Define the {@link ClassMatch} for filtering class.

apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/plugin/EnhanceContext.java

+package org.skywalking.apm.agent.core.plugin;
+
+/**
+ * The <code>EnhanceContext</code> represents the context or status for processing a class.
+ *
+ * Based on this context, the plugin core {@link org.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassEnhancePluginDefine}
+ * knows how to process the specific steps for every particular plugin.
+ *
+ * @author wusheng
+ */
+public class EnhanceContext {
+    private boolean isEnhanced = false;
+    /**
+     * The object has already been enhanced or extended.
+     * e.g. added the new field, or implemented the new interface
+     */
+    private boolean objectExtended = false;
+
+    public boolean isEnhanced() {
+        return isEnhanced;
+    }
+
+    public void initializationStageCompleted() {
+        isEnhanced = true;
+    }
+
+    public boolean isObjectExtended() {
+        return objectExtended;
+    }
+
+    public void extendObjectCompleted() {
+        objectExtended = true;
+    }
+}

apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/plugin/PluginFinder.java

@@ -22,7 +22,7 @@ import static net.bytebuddy.matcher.ElementMatchers.not;
  * @author wusheng
  */
 public class PluginFinder {
-    private final Map<String, AbstractClassEnhancePluginDefine> nameMatchDefine = new HashMap<String, AbstractClassEnhancePluginDefine>();
+    private final Map<String, LinkedList<AbstractClassEnhancePluginDefine>> nameMatchDefine = new HashMap<String, LinkedList<AbstractClassEnhancePluginDefine>>();
     private final List<AbstractClassEnhancePluginDefine> signatureMatchDefine = new LinkedList<AbstractClassEnhancePluginDefine>();
 
     public PluginFinder(List<AbstractClassEnhancePluginDefine> plugins) {
@@ -35,28 +35,34 @@ public class PluginFinder {
 
             if (match instanceof NameMatch) {
                 NameMatch nameMatch = (NameMatch)match;
-                nameMatchDefine.put(nameMatch.getClassName(), plugin);
+                LinkedList<AbstractClassEnhancePluginDefine> pluginDefines = nameMatchDefine.get(nameMatch.getClassName());
+                if (pluginDefines == null) {
+                    pluginDefines = new LinkedList<AbstractClassEnhancePluginDefine>();
+                    nameMatchDefine.put(nameMatch.getClassName(), pluginDefines);
+                }
+                pluginDefines.add(plugin);
             } else {
                 signatureMatchDefine.add(plugin);
             }
         }
     }
 
-    public AbstractClassEnhancePluginDefine find(TypeDescription typeDescription,
+    public List<AbstractClassEnhancePluginDefine> find(TypeDescription typeDescription,
         ClassLoader classLoader) {
+        List<AbstractClassEnhancePluginDefine> matchedPlugins = new LinkedList<AbstractClassEnhancePluginDefine>();
         String typeName = typeDescription.getTypeName();
         if (nameMatchDefine.containsKey(typeName)) {
-            return nameMatchDefine.get(typeName);
+            matchedPlugins.addAll(nameMatchDefine.get(typeName));
         }
 
         for (AbstractClassEnhancePluginDefine pluginDefine : signatureMatchDefine) {
             IndirectMatch match = (IndirectMatch)pluginDefine.enhanceClass();
             if (match.isMatch(typeDescription)) {
-                return pluginDefine;
+                matchedPlugins.add(pluginDefine);
             }
         }
 
-        return null;
+        return matchedPlugins;
     }
 
     public ElementMatcher<? super TypeDescription> buildMatch() {

apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/plugin/interceptor/enhance/ClassEnhancePluginDefine.java

@@ -6,6 +6,7 @@ import net.bytebuddy.implementation.MethodDelegation;
 import net.bytebuddy.implementation.SuperMethodCall;
 import net.bytebuddy.implementation.bind.annotation.Morph;
 import org.skywalking.apm.agent.core.plugin.AbstractClassEnhancePluginDefine;
+import org.skywalking.apm.agent.core.plugin.EnhanceContext;
 import org.skywalking.apm.agent.core.plugin.PluginException;
 import org.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint;
 import org.skywalking.apm.agent.core.plugin.interceptor.EnhanceException;
@@ -46,10 +47,11 @@ public abstract class ClassEnhancePluginDefine extends AbstractClassEnhancePlugi
      */
     @Override
     protected DynamicType.Builder<?> enhance(String enhanceOriginClassName,
-        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader) throws PluginException {
+        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader,
+        EnhanceContext context) throws PluginException {
         newClassBuilder = this.enhanceClass(enhanceOriginClassName, newClassBuilder, classLoader);
 
-        newClassBuilder = this.enhanceInstance(enhanceOriginClassName, newClassBuilder, classLoader);
+        newClassBuilder = this.enhanceInstance(enhanceOriginClassName, newClassBuilder, classLoader, context);
 
         return newClassBuilder;
     }
@@ -62,7 +64,8 @@ public abstract class ClassEnhancePluginDefine extends AbstractClassEnhancePlugi
      * @return new byte-buddy's builder for further manipulation.
      */
     private DynamicType.Builder<?> enhanceInstance(String enhanceOriginClassName,
-        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader) throws PluginException {
+        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader,
+        EnhanceContext context) throws PluginException {
         ConstructorInterceptPoint[] constructorInterceptPoints = getConstructorsInterceptPoints();
         InstanceMethodsInterceptPoint[] instanceMethodsInterceptPoints = getInstanceMethodsInterceptPoints();
 
@@ -83,16 +86,21 @@ public abstract class ClassEnhancePluginDefine extends AbstractClassEnhancePlugi
         }
 
         /**
-         * alter class source code.<br/>
+         * Manipulate class source code.<br/>
          *
          * new class need:<br/>
          * 1.Add field, name {@link #CONTEXT_ATTR_NAME}.
          * 2.Add a field accessor for this field.
          *
+         * And make sure the source codes manipulation only occurs once.
+         *
          */
-        newClassBuilder = newClassBuilder.defineField(CONTEXT_ATTR_NAME, Object.class, ACC_PRIVATE)
-            .implement(EnhancedInstance.class)
-            .intercept(FieldAccessor.ofField(CONTEXT_ATTR_NAME));
+        if (!context.isObjectExtended()) {
+            newClassBuilder = newClassBuilder.defineField(CONTEXT_ATTR_NAME, Object.class, ACC_PRIVATE)
+                .implement(EnhancedInstance.class)
+                .intercept(FieldAccessor.ofField(CONTEXT_ATTR_NAME));
+            context.extendObjectCompleted();
+        }
 
         /**
          * 2. enhance constructors

apm-sniffer/apm-agent/src/main/java/org/skywalking/apm/agent/SkyWalkingAgent.java

 package org.skywalking.apm.agent;
 
 import java.lang.instrument.Instrumentation;
+import java.util.List;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
@@ -9,6 +10,7 @@ import org.skywalking.apm.agent.core.boot.ServiceManager;
 import org.skywalking.apm.agent.core.conf.SnifferConfigInitializer;
 import org.skywalking.apm.agent.core.logging.EasyLogResolver;
 import org.skywalking.apm.agent.core.plugin.AbstractClassEnhancePluginDefine;
+import org.skywalking.apm.agent.core.plugin.EnhanceContext;
 import org.skywalking.apm.agent.core.plugin.PluginBootstrap;
 import org.skywalking.apm.agent.core.plugin.PluginException;
 import org.skywalking.apm.agent.core.plugin.PluginFinder;
@@ -54,13 +56,21 @@ public class SkyWalkingAgent {
             @Override
             public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription,
                 ClassLoader classLoader, JavaModule module) {
-                AbstractClassEnhancePluginDefine pluginDefine = pluginFinder.find(typeDescription, classLoader);
-                if (pluginDefine != null) {
-                    DynamicType.Builder<?> newBuilder = pluginDefine.define(typeDescription.getTypeName(), builder, classLoader);
-                    if (newBuilder != null) {
+                List<AbstractClassEnhancePluginDefine> pluginDefines = pluginFinder.find(typeDescription, classLoader);
+                if (pluginDefines.size() > 0) {
+                    DynamicType.Builder<?> newBuilder = builder;
+                    EnhanceContext context = new EnhanceContext();
+                    for (AbstractClassEnhancePluginDefine define : pluginDefines) {
+                        DynamicType.Builder<?> possibleNewBuilder = define.define(typeDescription.getTypeName(), newBuilder, classLoader, context);
+                        if (possibleNewBuilder != null) {
+                            newBuilder = possibleNewBuilder;
+                        }
+                    }
+                    if (context.isEnhanced()) {
                         logger.debug("Finish the prepare stage for {}.", typeDescription.getName());
-                        return newBuilder;
                     }
+
+                    return newBuilder;
                 }
 
                 logger.debug("Matched class {}, but ignore by finding mechanism.", typeDescription.getTypeName());