Agent: Optimize PluginBootService startup (apache#23113)

* Optimize plugin startup * add MetaDataContextsFactoryAdvice for metric * Adjust JDBCBackendTransactionManager to BackendTransactionManager * fix it error * add blank line * Optimize field names * optimize code * remove logback.xml from agent distribution * fix error of advisor config load * optimize code * format code
Byeongcheol-Kim · Dec 29, 2022 · ee98d50 · ee98d50
1 parent f230732
commit ee98d50
Show file tree

Hide file tree

Showing 12 changed files with 118 additions and 205 deletions.
diff --git a/...ootstrap/src/main/java/org/apache/shardingsphere/agent/bootstrap/ShardingSphereAgent.java b/...ootstrap/src/main/java/org/apache/shardingsphere/agent/bootstrap/ShardingSphereAgent.java
@@ -23,17 +23,15 @@
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.dynamic.scaffold.TypeValidation;
 import net.bytebuddy.matcher.ElementMatchers;
-import org.apache.shardingsphere.agent.config.advisor.AdvisorConfiguration;
-import org.apache.shardingsphere.agent.config.plugin.PluginConfiguration;
-import org.apache.shardingsphere.agent.bootstrap.classloader.AgentClassLoader;
 import org.apache.shardingsphere.agent.bootstrap.config.loader.PluginConfigurationLoader;
 import org.apache.shardingsphere.agent.bootstrap.logging.LoggingListener;
-import org.apache.shardingsphere.agent.bootstrap.plugin.PluginBootServiceManager;
 import org.apache.shardingsphere.agent.bootstrap.plugin.PluginJar;
 import org.apache.shardingsphere.agent.bootstrap.plugin.loader.AdvisorConfigurationLoader;
 import org.apache.shardingsphere.agent.bootstrap.plugin.loader.AgentPluginLoader;
 import org.apache.shardingsphere.agent.bootstrap.transformer.AgentJunction;
 import org.apache.shardingsphere.agent.bootstrap.transformer.AgentTransformer;
+import org.apache.shardingsphere.agent.config.advisor.AdvisorConfiguration;
+import org.apache.shardingsphere.agent.config.plugin.PluginConfiguration;
 
 import java.io.IOException;
 import java.lang.instrument.Instrumentation;
@@ -59,9 +57,6 @@ public static void premain(final String args, final Instrumentation instrumentat
         Collection<PluginJar> pluginJars = AgentPluginLoader.load();
         Map<String, AdvisorConfiguration> advisorConfigs = AdvisorConfigurationLoader.load(pluginJars, pluginConfigs.keySet(), isEnhancedForProxy);
         setUpAgentBuilder(instrumentation, pluginConfigs, pluginJars, advisorConfigs, isEnhancedForProxy);
-        if (isEnhancedForProxy) {
-            setupPluginBootService(pluginConfigs, pluginJars);
-        }
     }
 
     private static boolean isEnhancedForProxy() {
@@ -83,9 +78,4 @@ private static void setUpAgentBuilder(final Instrumentation instrumentation, fin
                 .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                 .with(new LoggingListener()).installOn(instrumentation);
     }
-
-    private static void setupPluginBootService(final Map<String, PluginConfiguration> pluginConfigs, final Collection<PluginJar> pluginJars) {
-        PluginBootServiceManager.startAllServices(pluginConfigs, AgentClassLoader.getClassLoader(), true);
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> PluginBootServiceManager.closeAllServices(pluginJars)));
-    }
 }
diff --git a/...er/builder/advise/ProxyAdviceFactory.java → ...strap/classloader/ClassLoaderContext.java b/...er/builder/advise/ProxyAdviceFactory.java → ...strap/classloader/ClassLoaderContext.java
@@ -15,36 +15,35 @@
  * limitations under the License.
  */
 
-package org.apache.shardingsphere.agent.bootstrap.transformer.builder.advise;
+package org.apache.shardingsphere.agent.bootstrap.classloader;
 
+import lombok.Getter;
 import lombok.RequiredArgsConstructor;
-import lombok.SneakyThrows;
-import org.apache.shardingsphere.agent.advice.AgentAdvice;
-import org.apache.shardingsphere.agent.bootstrap.classloader.AgentClassLoader;
+import org.apache.shardingsphere.agent.bootstrap.plugin.PluginJar;
 
+import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
- * Proxy Advice factory.
+ * Class loader context.
  */
 @RequiredArgsConstructor
-public final class ProxyAdviceFactory {
+@Getter
+public final class ClassLoaderContext {
 
-    private static final Map<String, AgentAdvice> CACHED_ADVICES = new ConcurrentHashMap<>();
+    private static final Map<ClassLoader, AgentClassLoader> AGENT_CLASS_LOADERS = new ConcurrentHashMap<>();
+
+    private final ClassLoader appClassLoader;
+
+    private final Collection<PluginJar> pluginJars;
 
     /**
-     * Get advice.
+     * Get agent class loader.
      *
-     * @param adviceClassName advice class name
-     * @return got advance
+     * @return agent class loader
      */
-    public AgentAdvice getAdvice(final String adviceClassName) {
-        return CACHED_ADVICES.computeIfAbsent(adviceClassName, this::createAdviceForProxy);
-    }
-
-    @SneakyThrows(ReflectiveOperationException.class)
-    private AgentAdvice createAdviceForProxy(final String adviceClassName) {
-        return (AgentAdvice) Class.forName(adviceClassName, true, AgentClassLoader.getClassLoader()).getDeclaredConstructor().newInstance();
+    public AgentClassLoader getAgentClassLoader() {
+        return AGENT_CLASS_LOADERS.computeIfAbsent(appClassLoader, key -> new AgentClassLoader(key, pluginJars));
     }
 }
diff --git a/...a/org/apache/shardingsphere/agent/bootstrap/plugin/loader/AdvisorConfigurationLoader.java b/...a/org/apache/shardingsphere/agent/bootstrap/plugin/loader/AdvisorConfigurationLoader.java
@@ -21,26 +21,26 @@
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
 import org.apache.shardingsphere.agent.bootstrap.classloader.AgentClassLoader;
+import org.apache.shardingsphere.agent.bootstrap.logging.LoggerFactory;
+import org.apache.shardingsphere.agent.bootstrap.logging.LoggerFactory.Logger;
 import org.apache.shardingsphere.agent.bootstrap.plugin.PluginJar;
 import org.apache.shardingsphere.agent.bootstrap.plugin.yaml.loader.YamlAdvisorsConfigurationLoader;
 import org.apache.shardingsphere.agent.bootstrap.plugin.yaml.swapper.YamlAdvisorsConfigurationSwapper;
-import org.apache.shardingsphere.agent.bootstrap.spi.PluginBootServiceRegistry;
 import org.apache.shardingsphere.agent.config.advisor.AdvisorConfiguration;
-import org.apache.shardingsphere.agent.spi.PluginBootService;
 
 import java.io.InputStream;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 
 /**
  * Advisor configuration loader.
  */
 @NoArgsConstructor(access = AccessLevel.PRIVATE)
 public final class AdvisorConfigurationLoader {
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(AgentPluginLoader.class);
+
     /**
      * Load advisor configurations.
      * 
@@ -52,22 +52,30 @@ public final class AdvisorConfigurationLoader {
     public static Map<String, AdvisorConfiguration> load(final Collection<PluginJar> pluginJars, final Collection<String> pluginTypes, final boolean isEnhancedForProxy) {
         Map<String, AdvisorConfiguration> result = new HashMap<>();
         AgentClassLoader.init(pluginJars);
-        for (PluginBootService each : PluginBootServiceRegistry.newInstances(AgentClassLoader.getClassLoader())) {
-            if (pluginTypes.contains(each.getType())) {
-                Collection<AdvisorConfiguration> advisorConfigs = YamlAdvisorsConfigurationSwapper
-                        .swapToObject(YamlAdvisorsConfigurationLoader.load(getAdvisorsResourceStream(each, isEnhancedForProxy)), each.getType());
-                result.putAll(advisorConfigs.stream().collect(Collectors.toMap(AdvisorConfiguration::getTargetClassName, Function.identity())));
+        for (String each : pluginTypes) {
+            InputStream advisorsResourceStream = getAdvisorsResourceStream(each, isEnhancedForProxy);
+            if (null == advisorsResourceStream) {
+                LOGGER.error("No configuration of advisor for type `{}`", each);
+                continue;
             }
+            Collection<AdvisorConfiguration> advisorConfigs = YamlAdvisorsConfigurationSwapper.swapToObject(YamlAdvisorsConfigurationLoader.load(advisorsResourceStream), each);
+            mergeAdvisorConfigurations(result, advisorConfigs);
         }
         return ImmutableMap.<String, AdvisorConfiguration>builder().putAll(result).build();
     }
 
-    private static InputStream getAdvisorsResourceStream(final PluginBootService pluginBootService, final boolean isEnhancedForProxy) {
-        InputStream result = pluginBootService.getClass().getResourceAsStream(getAdvisorsResourceFile(pluginBootService, (isEnhancedForProxy ? "proxy" : "jdbc") + "-advisors.yaml"));
-        return null == result ? pluginBootService.getClass().getResourceAsStream(getAdvisorsResourceFile(pluginBootService, "advisors.yaml")) : result;
+    private static InputStream getAdvisorsResourceStream(final String type, final boolean isEnhancedForProxy) {
+        InputStream result = AgentClassLoader.getClassLoader().getResourceAsStream(getAdvisorsResourceFile(type, (isEnhancedForProxy ? "proxy" : "jdbc") + "-advisors.yaml"));
+        return null == result ? AgentClassLoader.getClassLoader().getResourceAsStream(getAdvisorsResourceFile(type, "advisors.yaml")) : result;
+    }
+
+    private static String getAdvisorsResourceFile(final String type, final String fileName) {
+        return String.join("/", type.toLowerCase(), fileName);
     }
 
-    private static String getAdvisorsResourceFile(final PluginBootService pluginBootService, final String fileName) {
-        return String.join("/", "", pluginBootService.getType().toLowerCase(), fileName);
+    private static void mergeAdvisorConfigurations(final Map<String, AdvisorConfiguration> advisorConfigMap, final Collection<AdvisorConfiguration> advisorConfigs) {
+        for (AdvisorConfiguration each : advisorConfigs) {
+            advisorConfigMap.computeIfAbsent(each.getTargetClassName(), key -> new AdvisorConfiguration(each.getTargetClassName())).getAdvisors().addAll(each.getAdvisors());
+        }
     }
 }
diff --git a/...rc/main/java/org/apache/shardingsphere/agent/bootstrap/spi/PluginBootServiceRegistry.java b/...rc/main/java/org/apache/shardingsphere/agent/bootstrap/spi/PluginBootServiceRegistry.java
@@ -22,9 +22,7 @@
 import org.apache.shardingsphere.agent.spi.PluginBootService;
 
 import java.util.Collection;
-import java.util.LinkedList;
 import java.util.Optional;
-import java.util.ServiceLoader;
 
 /**
  *  Plugin boot service registry.
@@ -50,18 +48,4 @@ public static Optional<PluginBootService> getRegisteredService(final String type
     public static Collection<PluginBootService> getAllRegisteredServices() {
         return AgentServiceLoader.getServiceLoader(PluginBootService.class).getServices();
     }
-
-    /**
-     * Create new instances.
-     *
-     * @param classLoader class loader
-     * @return created instances
-     */
-    public static Collection<PluginBootService> newInstances(final ClassLoader classLoader) {
-        Collection<PluginBootService> result = new LinkedList<>();
-        for (PluginBootService each : ServiceLoader.load(PluginBootService.class, classLoader)) {
-            result.add(each);
-        }
-        return result;
-    }
 }
diff --git a/...src/main/java/org/apache/shardingsphere/agent/bootstrap/transformer/AgentTransformer.java b/...src/main/java/org/apache/shardingsphere/agent/bootstrap/transformer/AgentTransformer.java
@@ -25,14 +25,17 @@
 import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.utility.JavaModule;
 import org.apache.shardingsphere.agent.advice.TargetAdviceObject;
+import org.apache.shardingsphere.agent.bootstrap.classloader.ClassLoaderContext;
+import org.apache.shardingsphere.agent.bootstrap.plugin.PluginBootServiceManager;
 import org.apache.shardingsphere.agent.bootstrap.plugin.PluginJar;
-import org.apache.shardingsphere.agent.config.advisor.AdvisorConfiguration;
-import org.apache.shardingsphere.agent.config.plugin.PluginConfiguration;
 import org.apache.shardingsphere.agent.bootstrap.transformer.builder.MethodAdvisorBuilder;
 import org.apache.shardingsphere.agent.bootstrap.transformer.builder.advise.AdviceFactory;
+import org.apache.shardingsphere.agent.config.advisor.AdvisorConfiguration;
+import org.apache.shardingsphere.agent.config.plugin.PluginConfiguration;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Agent transformer.
@@ -42,6 +45,8 @@ public final class AgentTransformer implements Transformer {
 
     private static final String EXTRA_DATA = "_$EXTRA_DATA$_";
 
+    private static final AtomicBoolean STARTED_FLAG = new AtomicBoolean(false);
+
     private final Map<String, PluginConfiguration> pluginConfigs;
 
     private final Collection<PluginJar> pluginJars;
@@ -56,10 +61,16 @@ public Builder<?> transform(final Builder<?> builder, final TypeDescription type
         if (!advisorConfigs.containsKey(typeDescription.getTypeName())) {
             return builder;
         }
-        Builder<?> result = builder.defineField(EXTRA_DATA, Object.class, Opcodes.ACC_PRIVATE | Opcodes.ACC_VOLATILE).implement(TargetAdviceObject.class).intercept(FieldAccessor.ofField(EXTRA_DATA));
-        AdviceFactory adviceFactory = new AdviceFactory(classLoader, pluginConfigs, pluginJars, isEnhancedForProxy);
-        AdvisorConfiguration advisorConfig = advisorConfigs.get(typeDescription.getTypeName());
-        result = new MethodAdvisorBuilder(adviceFactory, advisorConfig, typeDescription).build(result);
-        return result;
+        ClassLoaderContext classLoaderContext = new ClassLoaderContext(classLoader, pluginJars);
+        startAllServices(classLoaderContext.getAgentClassLoader());
+        return new MethodAdvisorBuilder(new AdviceFactory(classLoaderContext), advisorConfigs.get(typeDescription.getTypeName()), typeDescription)
+                .build(builder.defineField(EXTRA_DATA, Object.class, Opcodes.ACC_PRIVATE | Opcodes.ACC_VOLATILE).implement(TargetAdviceObject.class).intercept(FieldAccessor.ofField(EXTRA_DATA)));
+    }
+
+    private void startAllServices(final ClassLoader classLoader) {
+        if (STARTED_FLAG.compareAndSet(false, true)) {
+            PluginBootServiceManager.startAllServices(pluginConfigs, classLoader, isEnhancedForProxy);
+            Runtime.getRuntime().addShutdownHook(new Thread(() -> PluginBootServiceManager.closeAllServices(pluginJars)));
+        }
     }
 }
diff --git a/...a/org/apache/shardingsphere/agent/bootstrap/transformer/builder/advise/AdviceFactory.java b/...a/org/apache/shardingsphere/agent/bootstrap/transformer/builder/advise/AdviceFactory.java
@@ -17,29 +17,23 @@
 
 package org.apache.shardingsphere.agent.bootstrap.transformer.builder.advise;
 
+import lombok.RequiredArgsConstructor;
+import lombok.SneakyThrows;
 import org.apache.shardingsphere.agent.advice.AgentAdvice;
-import org.apache.shardingsphere.agent.bootstrap.plugin.PluginJar;
-import org.apache.shardingsphere.agent.config.plugin.PluginConfiguration;
+import org.apache.shardingsphere.agent.bootstrap.classloader.ClassLoaderContext;
 
-import java.util.Collection;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Advice factory.
  */
+@RequiredArgsConstructor
 public final class AdviceFactory {
 
-    private final ProxyAdviceFactory proxyAdviceFactory;
+    private static final Map<String, AgentAdvice> CACHED_ADVICES = new ConcurrentHashMap<>();
 
-    private final JDBCAdviceFactory jdbcAdviceFactory;
-
-    private final boolean isEnhancedForProxy;
-
-    public AdviceFactory(final ClassLoader classLoader, final Map<String, PluginConfiguration> pluginConfigs, final Collection<PluginJar> pluginJars, final boolean isEnhancedForProxy) {
-        proxyAdviceFactory = new ProxyAdviceFactory();
-        jdbcAdviceFactory = new JDBCAdviceFactory(classLoader, pluginConfigs, pluginJars);
-        this.isEnhancedForProxy = isEnhancedForProxy;
-    }
+    private final ClassLoaderContext classLoaderContext;
 
     /**
      * Get advice.
@@ -48,6 +42,13 @@ public AdviceFactory(final ClassLoader classLoader, final Map<String, PluginConf
      * @return got advance
      */
     public AgentAdvice getAdvice(final String adviceClassName) {
-        return isEnhancedForProxy ? proxyAdviceFactory.getAdvice(adviceClassName) : jdbcAdviceFactory.getAdvice(adviceClassName);
+        String adviceInstanceCacheKey = String.format("%s_%s@%s", adviceClassName, classLoaderContext.getAppClassLoader().getClass().getName(),
+                Integer.toHexString(classLoaderContext.getAppClassLoader().hashCode()));
+        return CACHED_ADVICES.computeIfAbsent(adviceInstanceCacheKey, key -> createAdvice(adviceClassName));
+    }
+
+    @SneakyThrows(ReflectiveOperationException.class)
+    private AgentAdvice createAdvice(final String adviceClassName) {
+        return (AgentAdvice) Class.forName(adviceClassName, true, classLoaderContext.getAgentClassLoader()).getDeclaredConstructor().newInstance();
     }
 }
diff --git a/...g/apache/shardingsphere/agent/bootstrap/transformer/builder/advise/JDBCAdviceFactory.java b/...g/apache/shardingsphere/agent/bootstrap/transformer/builder/advise/JDBCAdviceFactory.java
diff --git a/agent/distribution/src/main/assembly/shardingsphere-agent-binary-distribution.xml b/agent/distribution/src/main/assembly/shardingsphere-agent-binary-distribution.xml
@@ -31,15 +31,6 @@
             <outputDirectory>conf</outputDirectory>
             <fileMode>0644</fileMode>
         </fileSet>
-        <fileSet>
-            <directory>src/main/resources</directory>
-            <includes>
-                <include>logback.xml</include>
-            </includes>
-            <outputDirectory>conf</outputDirectory>
-            <fileMode>0644</fileMode>
-        </fileSet>
-
         <fileSet>
             <directory>../../</directory>
             <includes>